diff --git a/.buildkite/README.md b/.buildkite/README.md index 033d112b8..11876bf03 100644 --- a/.buildkite/README.md +++ b/.buildkite/README.md @@ -8,15 +8,18 @@ build (PR and push to `main`): | Step | Purpose | Queue (machine) | |---|---|---| | `pre-commit` | pre-commit gate | `small_cpu_queue_premerge` (r6in.large) | -| `plugin-contracts` | plugin contract tests (19 files) | `medium_cpu_queue_premerge` (r6in.4xlarge) | -| `agent-adapter` | agent adapter tests (3 files) | `small_cpu_queue_premerge` | +| `plugin-contracts` | plugin contracts and CPU tests (27 files) | `medium_cpu_queue_premerge` (r6in.4xlarge) | +| `agent-adapter` | agent adapter tests (4 files) | `small_cpu_queue_premerge` | +| `upstream-sync-cpu` | mechanically synchronized upstream CPU tests | `medium_cpu_queue_premerge` | | `utils` | utils tests (`pytest tests/utils`) | `medium_cpu_queue_premerge` | -The three test steps depend on the pre-commit gate. Each suite runs its files +The four test steps depend on the pre-commit gate. Each suite runs its files sequentially inside one step because these queues boot a fresh EC2 instance -per job — a per-file matrix would be mostly boot + pip-install time. All -always-on CPU steps use the standard `python:3.11` image and install their -lightweight dependencies at runtime. +per job — a per-file matrix would be mostly boot + pip-install time. +Most always-on CPU steps use the standard `python:3.11` image and install their +lightweight dependencies at runtime. `upstream-sync-cpu` uses `VIME_CI_IMAGE` +(defaulting to `vllm/vime:latest`) because the synchronized GLM and checkpoint +tests import the image-pinned Megatron stack even though they do not allocate a GPU. ## Creating the pipeline (one-time, Buildkite UI) @@ -49,7 +52,7 @@ No secrets are required for these steps (WANDB etc. is GPU-suite only). The GPU suites are behind a **block step** (`:rocket: Run GPU test suites?`): click it in the Buildkite UI, multi-select the suites (`short`, -`vllm-config`, `megatron`, `precision`, `ckpt`), and the follow-up step +`vllm-config`, `megatron`, `vime-customized`, `precision`, `ckpt`), and the follow-up step generates one job per test via [`gpu_suites.py`](./gpu_suites.py) — the same `gpu_lock_exec.py` + `docker run` invocations used by the GPU jobs, including the per-test `VIME_TEST_USE_DEEPEP` / `VIME_TEST_USE_FP8_ROLLOUT` / @@ -61,12 +64,15 @@ reports a passing commit status even if nobody unblocks the GPU gate. GPU jobs run on the shared **`mithril-h100-pool`** queue, following the same pattern vllm-omni uses for it: each job is a Kubernetes pod (agent-stack-k8s `kubernetes` plugin) on an H100 SXM node, with GPUs allocated via -`nvidia.com/gpu` limits (4 or 8), a memory-backed `/dev/shm`, and the node's +`nvidia.com/gpu` limits (1 to 8), a memory-backed `/dev/shm`, and the node's `/mnt/hf-cache` mounted as `HF_HOME`. vime tests `hf download` their models at startup, so a warm HF cache is all they need. `WANDB_API_KEY` is not wired up yet; runs report without wandb until it's added (e.g. as a k8s secret in the pod spec). +Set `VIME_CI_IMAGE` to an immutable candidate digest for image-backed jobs; +otherwise they use `vllm/vime:latest`. Do not update `latest` before merge. + ## Keeping it in sync The test lists live in `pipeline.yml` and `gpu_suites.py`; update both together diff --git a/.buildkite/gpu_suites.py b/.buildkite/gpu_suites.py index 849e99abc..243bc22f8 100644 --- a/.buildkite/gpu_suites.py +++ b/.buildkite/gpu_suites.py @@ -23,31 +23,16 @@ import subprocess GPU_QUEUE = "mithril-h100-pool" -CI_IMAGE = "inferactinc/public:vime-latest" +CI_IMAGE = os.environ.get("VIME_CI_IMAGE", "vllm/vime:latest") HF_CACHE_HOST_PATH = "/mnt/hf-cache" HF_HOME = "/root/.cache/huggingface" NODE_INSTANCE_TYPE = "gpu-h100-sxm" -# Known hardware-fit failures on the pool's 80 GB H100s — test-level issues, -# not pipeline ones (PR #239, builds #6/#7): -# * gsm8k_async_short: FIXED — max-tokens-per-gpu reduced 9216→2048 (peak -# 39.6 GB on H200, well within H100 80 GB). Root cause was Qwen3.5 248k -# vocab × 5 logits copies in calculate_log_probs_and_entropy. -# * parallel_check: cross-layout grad-norm invariance (TP4+per-token-loss) -# diverges ~12% on ~11% of rollout data (bimodal: most <1.5%, outliers -# 10-20%). Confirmed same behavior in slime — Megatron FP reduction-order -# non-invariance, not a vime bug. -# soft_fail keeps them running and visible (orange) without failing the build. -SOFT_FAIL_ON_H100 = { - "test_qwen3_0.6B_parallel_check.py", -} - # (test_file, num_gpus, extra_args, env overrides) SUITES = { "short": [ ("test_qwen3.5_0.8B_gsm8k_async_short.py", 4, "", {}), ("test_qwen3.5_0.8B_gsm8k_short.py", 4, "", {}), - ("test_qwen2.5_0.5B_ppo_critic_only_short.py", 4, "", {}), ("test_qwen2.5_0.5B_fully_async_short.py", 4, "", {}), ], "vllm-config": [ @@ -57,26 +42,48 @@ ("test_vllm_config_mixed_offload_ft.py", 8, "", {}), ], "megatron": [ + ("test_full_disk_weight_update.py", 4, "", {}), ("test_quick_start_glm4_9B.py", 8, "", {}), ("test_glm4.7_30B_A3B_pd_mooncake.py", 8, "", {}), - ("test_qwen3_30B_A3B.py", 8, "", {"USE_DEEPEP": "1", "USE_FP8_ROLLOUT": "1"}), + ( + "test_qwen3_30B_A3B.py", + 8, + "", + {"USE_DEEPEP": "1", "USE_FP8_ROLLOUT": "1"}, + ), ("test_qwen3.6_35B_A3B_pd_mooncake.py", 8, "", {"USE_DEEPEP": "1"}), ("test_qwen3_30B_A3B_r3.py", 8, "", {"USE_DEEPEP": "1", "USE_FP8_ROLLOUT": "1", "ENABLE_EVAL": "0"}), ("test_qwen3_30B_A3B_r3.py", 8, "", {"ENABLE_EVAL": "0"}), ("test_qwen3_4B_ppo.py", 8, "", {}), ("test_qwen3_4B_ppo_disaggregate.py", 8, "", {}), ("test_qwen3_4B_ppo_train_critic_only.py", 8, "", {}), + ("test_ppo_logprob_entropy_gpu.py", 2, "", {}), + ("test_release_train.py", 4, "", {}), ("test_qwen3_4B_streaming_partial_rollout.py", 8, "", {}), ("test_moonlight_16B_A3B.py", 8, "", {}), ("test_moonlight_16B_A3B_r3.py", 8, "", {"ENABLE_EVAL": "0"}), + ("test_mimo_7B_mtp_only_grad.py", 8, "", {}), ("test_qwen2.5_0.5B_debug_rollout_then_train.py", 8, "", {}), ("test_qwen2.5_0.5B_opd_vllm.py", 8, "", {}), + ("test_qwen2.5_0.5B_fanout_short.py", 4, "", {}), + ("test_qwen2.5_0.5B_debug_train_dump_e2e.py", 8, "", {}), + ("test_qwen3_4B_external_pd.py", 6, "", {"VIME_TEST_UPDATE_MODE": "delta"}), + ], + "vime-customized": [ + ("test_qwen2_5_0_5B_non_colocate_pp.py", 4, "", {}), + ("test_geo3k_vlm_multi_turn_e2e.py", 1, "", {}), + ("test_qwen2.5_vl_3B_ep_disaggregation.py", 3, "", {}), + ("test_moonlight_16B_A3B_non_colocate_nccl.py", 8, "", {}), + ("test_qwen3_5_0_8B_top_p_cp2.py", 4, "", {}), ], "precision": [ ("test_qwen3_0.6B_parallel_check.py", 8, "", {}), ], "ckpt": [ - ("test_qwen3_4B_ckpt.py", 8, "", {}), + ("test_qwen3_4B_ckpt.py", 8, "--save-optimizer gpu --load-optimizer gpu", {}), + ("test_qwen3_4B_ckpt.py", 8, "--save-optimizer gpu --load-optimizer cpu", {}), + ("test_qwen3_4B_ckpt.py", 8, "--save-optimizer cpu --load-optimizer cpu", {}), + ("test_qwen3_4B_ckpt.py", 8, "--save-optimizer cpu --load-optimizer gpu", {}), ("test_qwen3_4B_ckpt.py", 8, "--async-save", {}), ], } @@ -107,9 +114,10 @@ def gpu_step(suite: str, test_file: str, num_gpus: int, extra_args: str, env: di {"name": "VIME_TEST_USE_DEEPEP", "value": vime_flags.get("USE_DEEPEP", "0")}, {"name": "VIME_TEST_USE_FP8_ROLLOUT", "value": vime_flags.get("USE_FP8_ROLLOUT", "0")}, {"name": "VIME_TEST_ENABLE_EVAL", "value": vime_flags.get("ENABLE_EVAL", "1")}, + {"name": "NCCL_NVLS_ENABLE", "value": env.get("NCCL_NVLS_ENABLE", "0")}, ] # anything else in env is passed to the pod verbatim (e.g. allocator knobs) - pod_env += [{"name": k, "value": v} for k, v in env.items() if k not in vime_flags] + pod_env += [{"name": k, "value": v} for k, v in env.items() if k not in vime_flags and k != "NCCL_NVLS_ENABLE"] # Set a stable commit identifier for downstream tooling. command = "\n".join( [ @@ -130,7 +138,12 @@ def gpu_step(suite: str, test_file: str, num_gpus: int, extra_args: str, env: di "command": command, "agents": {"queue": GPU_QUEUE}, "timeout_in_minutes": 360, - "retry": {"automatic": [{"exit_status": -1, "limit": 2}]}, + "retry": { + "automatic": [ + {"exit_status": -1, "limit": 2}, + {"exit_status": 1, "limit": 2}, + ] + }, "plugins": [ { "kubernetes": { @@ -159,9 +172,6 @@ def gpu_step(suite: str, test_file: str, num_gpus: int, extra_args: str, env: di } ], } - if test_file in SOFT_FAIL_ON_H100: - step["soft_fail"] = True - step["label"] = ":warning: " + step["label"] return step diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py index 7bad9d2ce..97517f9dd 100644 --- a/.buildkite/npu_suites.py +++ b/.buildkite/npu_suites.py @@ -31,7 +31,6 @@ ("test_qwen3_4B_npu.py", "npu-8", "", {}), ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}), ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}), - ("test_qwen3.5_35B_A3B_npu.py", "npu-16", "", {}), ("test_glm4.7_30B_A3B_npu.py", "npu-16", "", {}), ], "nightly": [], diff --git a/.buildkite/pipeline-npu-image.yaml b/.buildkite/pipeline-npu-image.yaml index 8f259b598..22b9ca82e 100644 --- a/.buildkite/pipeline-npu-image.yaml +++ b/.buildkite/pipeline-npu-image.yaml @@ -78,7 +78,6 @@ steps: --local context=. \ --local dockerfile=./docker \ --opt filename=Dockerfile.npu \ - --opt build-arg:BASE_IMAGE=swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/vllm-ascend/vllm-ascend \ --opt build-arg:APTMIRROR=http://cache-service.nginx-pypi-cache.svc.cluster.local:8081 \ --opt build-arg:PIP_INDEX_URL=http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ --secret id=dockerconfig,src=/home/user/.docker/config.json \ diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 5ae2e5900..dbc8b2258 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,6 +26,10 @@ steps: automatic: - exit_status: -1 # agent lost (fresh instance failed to boot) limit: 2 + - exit_status: 1 # transient package-index failure + limit: 2 + - exit_status: 3 # pre-commit environment install failure + limit: 2 command: | docker run --rm \ -e GIT_CONFIG_PARAMETERS="'safe.directory=/workspace'" \ @@ -46,6 +50,8 @@ steps: automatic: - exit_status: -1 limit: 2 + - exit_status: 1 + limit: 2 command: | # GLOO/TP_SOCKET_IFNAME=lo: the torch.distributed tests (e.g. # test_metric_report_dist) rendezvous over localhost; inside a @@ -58,11 +64,13 @@ steps: python:3.11 bash -c ' set -euo pipefail pip install -q torch --index-url https://download.pytorch.org/whl/cpu - pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle + pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle blake3 xxhash zstandard psutil wandb pip install -q -e . --no-deps python tests/test_megatron_argument_validation.py python tests/test_value_temperature.py - python tests/test_rollout_validation.py + python tests/test_docs_consistency.py + python tests/test_placement_group.py + python tests/test_external_vllm_engines.py python tests/plugin_contracts/test_plugin_rollout_contracts.py python tests/plugin_contracts/test_plugin_runtime_hook_contracts.py python tests/plugin_contracts/test_plugin_path_loading_contracts.py @@ -78,6 +86,12 @@ steps: python tests/test_metric_report_dist.py python tests/test_loss_cp_invariance.py python tests/test_sample.py + python tests/test_cispo_loss.py + python tests/test_logprob_response_spans.py + python tests/test_empty_colocated_weight_bucket.py + python tests/test_expert_routing.py + python tests/test_reloadable_process_group_memory_check.py + python tests/test_ppo_logprob_entropy.py python tests/utils/test_hf_checkpoint_saver.py ' @@ -91,6 +105,8 @@ steps: automatic: - exit_status: -1 limit: 2 + - exit_status: 1 + limit: 2 command: | docker run --rm --shm-size=2g \ -e GIT_CONFIG_PARAMETERS="'safe.directory=/workspace'" \ @@ -98,12 +114,69 @@ steps: python:3.11 bash -c ' set -euo pipefail pip install -q torch --index-url https://download.pytorch.org/whl/cpu - pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle + pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle blake3 xxhash zstandard wandb pip install -q openai openai-agents anthropic pip install -q -e . --no-deps - python tests/test_agent_trajectory.py - python tests/test_agent_adapters.py - python tests/test_agent_sdk_adapters.py + python tests/test_agent/test_adapters.py + python tests/test_agent/test_harness.py + python tests/test_agent/test_trajectory_manager_branching.py + python tests/test_agent/test_agent_rollout_cpu.py + ' + + - label: ":pytest: synchronized upstream CPU tests" + key: upstream-sync-cpu + depends_on: pre-commit + agents: + queue: medium_cpu_queue_premerge + timeout_in_minutes: 45 + retry: + automatic: + - exit_status: -1 + limit: 2 + - exit_status: 1 + limit: 2 + command: | + docker run --rm --network host --ipc=host \ + -e GIT_CONFIG_PARAMETERS="'safe.directory=/workspace'" \ + -e GLOO_SOCKET_IFNAME=lo -e TP_SOCKET_IFNAME=lo \ + -v "$$PWD:/workspace" -w /workspace \ + "$${VIME_CI_IMAGE:-vllm/vime:latest}" bash -lc ' + set -euo pipefail + pip install -q -e . --no-deps --break-system-packages + for test_file in \ + tests/test_advantage_whiten_cp.py \ + tests/test_accelerator.py \ + tests/test_block_fp8_zero_block.py \ + tests/test_deep_ep_tms_patch.py \ + tests/test_discounted_returns.py \ + tests/test_eval_config.py \ + tests/test_filter_long_prompt.py \ + tests/test_fully_async_rollout.py \ + tests/test_glm52_layerwise_comparison.py \ + tests/test_glm5_indexer_q_norm.py \ + tests/test_glm5_indexer_short_context.py \ + tests/test_hf_to_megatron.py \ + tests/test_layerwise_alignment.py \ + tests/test_model_provider_freeze.py \ + tests/test_policy_loss.py \ + tests/test_ppo_kl_metric.py \ + tests/test_process_rollout_data.py \ + tests/test_qwen3_5_vl_native.py \ + tests/test_qwen3_linear_attention_cu_seqlens.py \ + tests/test_read_file_slicing.py \ + tests/test_reloadable_process_group_world.py \ + tests/test_rollout_data_utils.py \ + tests/test_rollout_metrics.py \ + tests/test_rollout_sample_hooks.py \ + tests/test_vllm_rollout.py \ + tests/test_agent/test_sandbox_exec_and_wait.py \ + tests/test_stateless_adam.py \ + tests/test_tau_bench_token_delta.py \ + tests/test_train_data_utils.py \ + tests/test_update_weight_factory.py \ + tests/observability/test_trace_utils.py; do + python -m pytest "$$test_file" + done ' - label: ":pytest: utils tests" @@ -116,6 +189,8 @@ steps: automatic: - exit_status: -1 limit: 2 + - exit_status: 1 + limit: 2 command: | docker run --rm --network host --ipc=host \ -e GIT_CONFIG_PARAMETERS="'safe.directory=/workspace'" \ @@ -123,7 +198,7 @@ steps: python:3.11 bash -lc ' set -euo pipefail pip install -q torch --index-url https://download.pytorch.org/whl/cpu - pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle + pip install -q pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors transformers cloudpickle blake3 xxhash zstandard pip install -q -e . --no-deps python -m pytest tests/utils ' @@ -143,15 +218,17 @@ steps: multiple: true required: true options: - - label: "run-ci-short — 4 GPU, 4 tests" + - label: "run-ci-short — 4 GPU, 3 tests" value: short - label: "run-ci-vllm-config — 8 GPU, 4 tests" value: vllm-config - - label: "run-ci-megatron — 8 GPU, 14 runs" + - label: "run-ci-megatron — up to 8 GPU, 21 runs" value: megatron + - label: "run-ci-vime-customized — 1–8 GPU, 6 tests" + value: vime-customized - label: "run-ci-precision — 8 GPU, 1 test" value: precision - - label: "run-ci-ckpt — 8 GPU, 2 runs" + - label: "run-ci-ckpt — 8 GPU, 5 runs" value: ckpt - label: ":pipeline: upload selected GPU suites" diff --git a/.claude/skills/add-tests-and-ci/SKILL.md b/.claude/skills/add-tests-and-ci/SKILL.md index 98b2140e6..a729a36d3 100644 --- a/.claude/skills/add-tests-and-ci/SKILL.md +++ b/.claude/skills/add-tests-and-ci/SKILL.md @@ -40,15 +40,33 @@ if __name__ == "__main__": - `run-ci-changed` extracts a top-level `NUM_GPUS = ` constant from added/modified `tests/test_*.py` and `tests/plugin_contracts/test_*.py`; if missing, it defaults to 8 GPUs. Set `NUM_GPUS = 0` for CPU-only tests. - For GPU/e2e tests, follow the nearby file pattern (`prepare()`, `execute()`, `NUM_GPUS`, and any model/dataset constants). -### Step 3: Run Local Validation +### Step 3: Register Tests in GitHub CI + +Whenever adding, moving, or renaming a test file, update the GitHub workflow template before finishing: + +1. Add the test to the appropriate matrix in `.github/workflows/pr-test.yml.j2`. + - CPU-only pytest/unit tests usually belong in `cpu-unittest` with `num_gpus: 0`. + - GPU/e2e tests should be placed beside the nearest similar model/path test with the matching `num_gpus` and environment fields. +2. Regenerate workflows: + +```bash +python .github/workflows/generate_github_workflows.py +``` + +3. Include both `.github/workflows/pr-test.yml.j2` and the generated `.github/workflows/pr-test.yml` in the change set. + +Only skip fixed matrix registration when the test is intentionally helper-only or manually invoked; state that reason in the final response. + +### Step 4: Run Local Validation - Run the exact existing test files you changed, if any. +- For new registered tests, run the same shape CI will use, for example `python tests/test_new_file.py`. - Run repository-wide checks only when they are already part of the task or workflow. - Avoid documenting placeholder test commands that may not exist in the current tree. -### Step 4: Update Workflow Template Correctly +### Step 5: Keep Workflow Template as Source of Truth -For CI workflow changes: +For CI workflow changes unrelated to a new, moved, or renamed test: 1. Edit `.github/workflows/pr-test.yml.j2` 2. Regenerate workflows: @@ -59,11 +77,12 @@ python .github/workflows/generate_github_workflows.py 3. Include both the template and generated workflow file in the change set (`.j2` and `.yml`). If the user asked for a commit, commit both. -### Step 5: Provide Verifiable PR Notes +### Step 6: Provide Verifiable PR Notes Include: - Which tests were added/changed +- Where each new/renamed test was registered in `.github/workflows/pr-test.yml.j2` - Exact commands executed - GPU assumptions for each test path - Why this coverage protects against regression @@ -71,6 +90,7 @@ Include: ## Common Mistakes - Editing generated workflow file only +- Relying on `run-ci-changed` discovery for a new test that should run in the regular PR matrix - Forgetting `NUM_GPUS = 0` on a CPU-only changed test, causing `run-ci-changed` to default to 8 GPUs - Adding a CPU pytest file that passes under `pytest tests/foo.py` but fails under CI's `python tests/foo.py` - Adding tests without following existing constants/conventions diff --git a/.claude/skills/vime-code-review-preferences/SKILL.md b/.claude/skills/vime-code-review-preferences/SKILL.md new file mode 100644 index 000000000..5f3be3fd1 --- /dev/null +++ b/.claude/skills/vime-code-review-preferences/SKILL.md @@ -0,0 +1,29 @@ +--- +name: vime-code-review-preferences +description: Use when reviewing or editing vime code, especially refactors around helper APIs, branch selection, argument validation, or recurring reviewer preferences about avoiding unnecessary wrappers and making control flow self-explanatory. +--- + +# Vime Code Review Preferences + +Apply these lightweight review heuristics when changing vime code. + +## Prefer Direct APIs Over Thin Wrappers + +- Remove helper layers that only rename a call, format one path, or forward arguments without owning meaningful behavior. +- Prefer calling the concrete reusable API directly, for example a `*_to_path` helper when the caller already knows the destination path. +- Keep a wrapper only if it owns a real boundary: compatibility, validation, nontrivial error policy, lifecycle management, metrics/logging semantics, async/retry behavior, or cross-module ownership. +- Avoid moving a redundant wrapper's body into another file just to preserve the wrapper shape. Inline the simple call at the natural ownership site. +- When removing a wrapper, search for sibling wrappers and nearby helpers with `rg` and delete confirmed dead functions in the same pass. +- Treat single-use convenience functions as suspicious when their only job is path formatting plus forwarding. Prefer the caller owning that one line. + +## Make Branches Explain Themselves + +- Order conditionals by semantic precedence: special transport/lifecycle modes first, then explicit mode choices, then default paths. +- Prefer predicates that fully describe the branch, such as `mode == "full" and transport == "disk"`, over a broad predicate followed by an assert that explains what the branch really meant. +- Use asserts as invariants for impossible states after validation, not as a substitute for clear branch conditions. + +## Keep Abstractions Honest + +- Add an abstraction only when it removes real duplication, hides fragile mechanics, or clarifies ownership. +- When a review comment points out repeated indirection, look for a smaller public surface rather than adding another alias. +- Preserve existing behavior intentionally. If cleanup changes error handling, logging, or failure visibility, call that out in the final response. diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 2996653dc..13a3caaf1 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -65,7 +65,7 @@ body: - vime version: - Python version: - PyTorch version: - - CUDA version: + - CUDA/ROCm version: - GPU type and count: - OS: - vLLM version: diff --git a/.github/ISSUE_TEMPLATE/question.yml b/.github/ISSUE_TEMPLATE/question.yml index ac116bfab..2c541e524 100644 --- a/.github/ISSUE_TEMPLATE/question.yml +++ b/.github/ISSUE_TEMPLATE/question.yml @@ -49,7 +49,7 @@ body: - vime version: - Python version: - PyTorch version: - - CUDA version: + - CUDA/ROCm version: - GPU type and count: - OS: diff --git a/.github/workflows/bot-slash-lint.yaml b/.github/workflows/bot-slash-lint.yaml deleted file mode 100644 index 85ae508c2..000000000 --- a/.github/workflows/bot-slash-lint.yaml +++ /dev/null @@ -1,110 +0,0 @@ -name: Slash Command Handler - -on: - issue_comment: - types: [created, edited] - -permissions: - contents: write # Required to push commits back to PR branch - actions: write # Required to rerun workflows - issues: write # Required for comment reactions in some contexts - -jobs: - slash_lint_codebase: - # Only run if it is a PR comment with a recognized command - if: > - github.event_name == 'issue_comment' && - github.event.issue.pull_request && - ( - contains(github.event.comment.body, '/tag-run-lint') || - contains(github.event.comment.body, '/run-lint') - ) - runs-on: ubuntu-latest - steps: - - name: React to command comment (ack) - if: always() - uses: actions/github-script@v8 - with: - script: | - const commentId = context.payload.comment.id; - // Add an eyes reaction to acknowledge the command - await github.request('POST /repos/{owner}/{repo}/issues/comments/{comment_id}/reactions', { - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: commentId, - content: 'eyes' - }); - - - name: Check out Git repository - uses: actions/checkout@v6 - with: - repository: ${{ github.repository }} - ref: refs/pull/${{ github.event.issue.number }}/head - token: ${{ secrets.GITHUB_TOKEN }} - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: '3.10' - - - name: Run pre-commit hooks - continue-on-error: true - uses: pre-commit/action@v3.0.1 - - - name: Get PR branch name - id: get_branch - run: | - BRANCH_NAME=$(gh pr view ${{ github.event.issue.number }} --json headRefName --jq '.headRefName') - echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - - - name: Check if there are any changes - id: verify_diff - run: | - git diff --quiet . || echo "changed=true" >> $GITHUB_OUTPUT - - - name: Commit files - if: steps.verify_diff.outputs.changed == 'true' - run: | - git config --local user.email "action@github.com" - git config --local user.name "GitHub Action" - git add . - git commit -m "[CI-Lint] Fix code style issues with pre-commit ${{ github.sha }}" -a - git push origin HEAD:refs/heads/${{ steps.get_branch.outputs.branch_name }} - - cleanup_reaction: - # Always run after the main job completes (success, failure, or cancelled) - if: always() - needs: slash_lint_codebase - runs-on: ubuntu-latest - steps: - - name: Remove initial ack reaction - uses: actions/github-script@v8 - with: - script: | - const commentId = context.payload.comment.id; - // List reactions on the comment - const reactions = await github.rest.reactions.listForIssueComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: commentId - }).then(r => r.data); - // Find the 'eyes' reaction added by this workflow bot - const target = reactions.find(r => r.content === 'eyes' && r.user && r.user.login === 'github-actions[bot]'); - if (target) { - try { - await github.rest.reactions.deleteForIssueComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: commentId, - reaction_id: target.id - }); - core.info(`Successfully deleted eyes reaction (${target.id})`); - } catch (err) { - // Non-fatal: reaction may already be gone or inaccessible - core.info(`Could not delete eyes reaction (${target.id}): ${err.message || err.status || 'unknown error'}`); - } - } else { - core.info('No eyes reaction from github-actions[bot] found to remove.'); - } diff --git a/.github/workflows/generate_github_workflows.py b/.github/workflows/generate_github_workflows.py deleted file mode 100644 index ee98bfcb5..000000000 --- a/.github/workflows/generate_github_workflows.py +++ /dev/null @@ -1,37 +0,0 @@ -from pathlib import Path -import jinja2 - - -def main(): - """ - Generates GitHub workflow YAML files from Jinja2 templates. - """ - workflows_dir = Path(__file__).parent - print(f"Scan dir: {workflows_dir}") - env = jinja2.Environment( - loader=jinja2.FileSystemLoader(str(workflows_dir)), - block_start_string="<%", - block_end_string="%>", - variable_start_string="<<", - variable_end_string=">>", - ) - - for template_path in workflows_dir.glob("*.yml.j2"): - template = env.get_template(template_path.name) - content = template.render() - - yaml_path = template_path.with_suffix("") - with open(yaml_path, "w") as f: - f.write( - "#" * 80 - + "\n# This file is auto-generated from the .j2 file via generate_github_workflows.py. Do not edit manually.\n" - + "#" * 80 - + "\n" - ) - f.write(content) - - print(f"Generated {yaml_path} from {template_path}") - - -if __name__ == "__main__": - main() diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml deleted file mode 100644 index a3f253aef..000000000 --- a/.github/workflows/pr-test.yml +++ /dev/null @@ -1,778 +0,0 @@ -################################################################################ -# This file is auto-generated from the .j2 file via generate_github_workflows.py. Do not edit manually. -################################################################################ - -name: PR Test - -on: - # Push to main triggers the default GitHub-hosted jobs (cheap CPU/unit and - # agent adapter checks), catching PR-pair regressions where two PRs pass - # individually but main is broken after both land. GPU jobs stay label-gated - # — see each job's `if:` below — so push events never burn the self-hosted - # fleet. - push: - branches: [main] - pull_request: - branches: [main] - types: [opened, reopened, synchronize, labeled] - workflow_dispatch: - inputs: - infinite_run: - description: 'Run training infinitely' - required: false - type: boolean - default: false - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - - - e2e-test-short: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-short')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - e2e-test-vllm-config: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-vllm-config')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_vllm_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_vllm_config_distributed.py"}, {"num_gpus": 8, "test_file": "test_vllm_config_mixed_offload.py"}, {"num_gpus": 8, "test_file": "test_vllm_config_mixed_offload_ft.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - e2e-test-megatron: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-megatron')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - e2e-test-precision: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-precision')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - e2e-test-ckpt: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-ckpt')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - cpu-unittest: - - if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' || github.event_name == 'push' - - - runs-on: ubuntu-latest - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_value_temperature.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}, {"num_gpus": 0, "test_file": "test_rm_deepscaler.py"}, {"num_gpus": 0, "test_file": "test_rm_f1.py"}, {"num_gpus": 0, "test_file": "test_rm_gpqa.py"}, {"num_gpus": 0, "test_file": "test_rm_math.py"}, {"num_gpus": 0, "test_file": "test_rm_math_dapo.py"}, {"num_gpus": 0, "test_file": "test_dp_schedule.py"}, {"num_gpus": 0, "test_file": "test_cp_utils.py"}, {"num_gpus": 0, "test_file": "test_metric_report.py"}, {"num_gpus": 0, "test_file": "test_metric_report_dist.py"}, {"num_gpus": 0, "test_file": "test_loss_cp_invariance.py"}, {"num_gpus": 0, "test_file": "test_sample.py"}, {"num_gpus": 0, "test_file": "utils/test_hf_checkpoint_saver.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - cache: 'pip' - - - name: Install dependencies - shell: bash - run: | - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors - - - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps - - - - name: Execute - shell: bash - run: | - - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - TEST_ARGS="${{ matrix.info.test_args || '' }}" - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - - - agent-adapter-test: - - if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' || github.event_name == 'push' - - - runs-on: ubuntu-latest - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 0, "test_file": "test_agent_trajectory.py"}, {"num_gpus": 0, "test_file": "test_agent_adapters.py"}, {"num_gpus": 0, "test_file": "test_agent_sdk_adapters.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - cache: 'pip' - - - name: Install dependencies - shell: bash - run: | - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors - - pip install openai openai-agents anthropic - - - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps - - - - name: Execute - shell: bash - run: | - - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - TEST_ARGS="${{ matrix.info.test_args || '' }}" - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - - - e2e-test-image: - - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-image')) - - - runs-on: self-hosted - - strategy: - fail-fast: false - matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - - - name: Execute - shell: bash - run: | - - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-test-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' - - - - e2e-test-changed-detect: - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-changed')) - runs-on: self-hosted - outputs: - matrix: ${{ steps.detect.outputs.matrix }} - has_tests: ${{ steps.detect.outputs.has_tests }} - steps: - - name: Checkout repository - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - - name: Detect changed tests - id: detect - shell: bash - run: | - CHANGED=$(git diff --name-only --diff-filter=AM origin/main...HEAD -- 'tests/test_*.py' 'tests/plugin_contracts/test_*.py' || true) - if [ -z "$CHANGED" ]; then - echo "No new or modified test files found." - echo "has_tests=false" >> $GITHUB_OUTPUT - echo 'matrix={"info":[]}' >> $GITHUB_OUTPUT - else - echo "Changed test files:" - echo "$CHANGED" - MATRIX="[" - FIRST=true - for filepath in $CHANGED; do - # Extract NUM_GPUS from the test file, default to 8 - NGPU=$(grep -oP '^NUM_GPUS\s*=\s*\K\d+' "$filepath" | head -1) - NGPU=${NGPU:-8} - if [ "$FIRST" = true ]; then FIRST=false; else MATRIX+=","; fi - MATRIX+="{\"test_file\":\"$filepath\",\"num_gpus\":$NGPU}" - done - MATRIX+="]" - echo "has_tests=true" >> $GITHUB_OUTPUT - echo "matrix={\"info\":$MATRIX}" >> $GITHUB_OUTPUT - echo "Generated matrix: $MATRIX" - fi - - e2e-test-changed: - needs: e2e-test-changed-detect - if: needs.e2e-test-changed-detect.outputs.has_tests == 'true' - runs-on: self-hosted - strategy: - fail-fast: false - matrix: ${{ fromJson(needs.e2e-test-changed-detect.outputs.matrix) }} - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - name: Execute - shell: bash - run: | - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' \ No newline at end of file diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 deleted file mode 100644 index 48a766f88..000000000 --- a/.github/workflows/pr-test.yml.j2 +++ /dev/null @@ -1,376 +0,0 @@ -<% set jobs = { - 'e2e-test-short': { - 'label': 'run-ci-short', - 'tests': [ - {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4}, - ], - }, - 'e2e-test-vllm-config': { - 'label': 'run-ci-vllm-config', - 'tests': [ - {'test_file': 'test_qwen2.5_0.5B_vllm_config.py', 'num_gpus': 8}, - {'test_file': 'test_qwen2.5_0.5B_vllm_config_distributed.py', 'num_gpus': 8}, - {'test_file': 'test_vllm_config_mixed_offload.py', 'num_gpus': 8}, - {'test_file': 'test_vllm_config_mixed_offload_ft.py', 'num_gpus': 8}, - ], - }, - 'e2e-test-megatron': { - 'label': 'run-ci-megatron', - 'tests': [ - {'test_file': 'test_quick_start_glm4_9B.py', 'num_gpus': 8}, - {'test_file': 'test_glm4.7_30B_A3B_pd_mooncake.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1'}, - {'test_file': 'test_qwen3.6_35B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'}, - {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'}, - {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'}, - {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ppo_disaggregate.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_streaming_partial_rollout.py', 'num_gpus': 8}, - {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, - {'test_file': 'test_moonlight_16B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'}, - {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, - {'test_file': 'test_qwen2.5_0.5B_opd_vllm.py', 'num_gpus': 8}, - ], - }, - 'e2e-test-precision': { - 'label': 'run-ci-precision', - 'tests': [ - {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8}, - ], - }, - 'e2e-test-ckpt': { - 'label': 'run-ci-ckpt', - 'tests': [ - {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8}, - ], - }, - - 'cpu-unittest': { - 'label': 'run-ci-cpu-unittest', - 'always': True, - 'cpu': True, - 'tests': [ - {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0}, - {'test_file': 'test_value_temperature.py', 'num_gpus': 0}, - {'test_file': 'test_rollout_validation.py', 'num_gpus': 0}, - {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0}, - {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0}, - {'test_file': 'plugin_contracts/test_plugin_path_loading_contracts.py', 'num_gpus': 0}, - {'test_file': 'plugin_contracts/test_plugin_generate_contracts.py', 'num_gpus': 0}, - {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0}, - {'test_file': 'test_rm_f1.py', 'num_gpus': 0}, - {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0}, - {'test_file': 'test_rm_math.py', 'num_gpus': 0}, - {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0}, - {'test_file': 'test_dp_schedule.py', 'num_gpus': 0}, - {'test_file': 'test_cp_utils.py', 'num_gpus': 0}, - {'test_file': 'test_metric_report.py', 'num_gpus': 0}, - {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0}, - {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0}, - {'test_file': 'test_sample.py', 'num_gpus': 0}, - {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0}, - ], - }, - - 'agent-adapter-test': { - 'label': 'run-ci-agent-adapter', - 'always': True, - 'cpu': True, - 'extra_pip_deps': 'openai openai-agents anthropic', - 'tests': [ - {'test_file': 'test_agent_trajectory.py', 'num_gpus': 0}, - {'test_file': 'test_agent_adapters.py', 'num_gpus': 0}, - {'test_file': 'test_agent_sdk_adapters.py', 'num_gpus': 0}, - ], - }, - - 'e2e-test-image': { - 'label': 'run-ci-image', - 'image': 'inferactinc/public:vime-test-latest', - 'tests': [ - {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4}, - {'test_file': 'test_quick_start_glm4_9B.py', 'num_gpus': 8}, - {'test_file': 'test_glm4.7_30B_A3B_pd_mooncake.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3.6_35B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'}, - {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, - {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8}, - {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, - {'test_file': 'test_qwen2.5_0.5B_opd_vllm.py', 'num_gpus': 8}, - ], - }, -} %> -name: PR Test - -on: - # Push to main triggers the default GitHub-hosted jobs (cheap CPU/unit and - # agent adapter checks), catching PR-pair regressions where two PRs pass - # individually but main is broken after both land. GPU jobs stay label-gated - # — see each job's `if:` below — so push events never burn the self-hosted - # fleet. - push: - branches: [main] - pull_request: - branches: [main] - types: [opened, reopened, synchronize, labeled] - workflow_dispatch: - inputs: - infinite_run: - description: 'Run training infinitely' - required: false - type: boolean - default: false - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - -<% for job_name, config in jobs.items() %> - << job_name >>: -<% if config.get('always') %> - if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' || github.event_name == 'push' -<% else %> - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, '<< config.label >>')) -<% endif %> -<% if config.get('cpu') %> - runs-on: ubuntu-latest -<% else %> - runs-on: self-hosted -<% endif %> - strategy: - fail-fast: false - matrix: - info: << config.tests | tojson >> - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 -<% if config.get('cpu') %> - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - cache: 'pip' - - - name: Install dependencies - shell: bash - run: | - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors -<% if config.get('extra_pip_deps') %> - pip install << config.extra_pip_deps >> -<% endif %> - - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps -<% else %> -<% endif %> - - - name: Execute - shell: bash - run: | -<% if config.get('cpu') %> - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - TEST_ARGS="${{ matrix.info.test_args || '' }}" - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi -<% else %> - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - << config.image if config.image else 'inferactinc/public:vime-latest' >> \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' -<% endif %> -<% endfor %> - - e2e-test-changed-detect: - if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-changed')) - runs-on: self-hosted - outputs: - matrix: ${{ steps.detect.outputs.matrix }} - has_tests: ${{ steps.detect.outputs.has_tests }} - steps: - - name: Checkout repository - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - - name: Detect changed tests - id: detect - shell: bash - run: | - CHANGED=$(git diff --name-only --diff-filter=AM origin/main...HEAD -- 'tests/test_*.py' 'tests/plugin_contracts/test_*.py' || true) - if [ -z "$CHANGED" ]; then - echo "No new or modified test files found." - echo "has_tests=false" >> $GITHUB_OUTPUT - echo 'matrix={"info":[]}' >> $GITHUB_OUTPUT - else - echo "Changed test files:" - echo "$CHANGED" - MATRIX="[" - FIRST=true - for filepath in $CHANGED; do - # Extract NUM_GPUS from the test file, default to 8 - NGPU=$(grep -oP '^NUM_GPUS\s*=\s*\K\d+' "$filepath" | head -1) - NGPU=${NGPU:-8} - if [ "$FIRST" = true ]; then FIRST=false; else MATRIX+=","; fi - MATRIX+="{\"test_file\":\"$filepath\",\"num_gpus\":$NGPU}" - done - MATRIX+="]" - echo "has_tests=true" >> $GITHUB_OUTPUT - echo "matrix={\"info\":$MATRIX}" >> $GITHUB_OUTPUT - echo "Generated matrix: $MATRIX" - fi - - e2e-test-changed: - needs: e2e-test-changed-detect - if: needs.e2e-test-changed-detect.outputs.has_tests == 'true' - runs-on: self-hosted - strategy: - fail-fast: false - matrix: ${{ fromJson(needs.e2e-test-changed-detect.outputs.matrix) }} - defaults: - run: - working-directory: ${{ github.workspace }} - env: - GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} - WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - - name: Execute - shell: bash - run: | - docker run --pull=always --rm \ - --privileged \ - --cap-add SYS_NICE \ - --security-opt seccomp=unconfined \ - --network host \ - --gpus all \ - --ipc=host \ - --shm-size=16g \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --memory=0 \ - --memory-swap=0 \ - -e http_proxy \ - -e https_proxy \ - -e HTTP_PROXY \ - -e HTTPS_PROXY \ - -e GITHUB_COMMIT_NAME \ - -e WANDB_API_KEY \ - -e VIME_TEST_ENABLE_INFINITE_RUN \ - -e VIME_TEST_USE_DEEPEP \ - -e VIME_TEST_USE_FP8_ROLLOUT \ - -e VIME_TEST_ENABLE_EVAL \ - -e TEST_FILE="${{ matrix.info.test_file }}" \ - -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ - -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ - -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ - -v /mnt/nvme0n1/vime_ci:/data/vime_ci \ - -v /mnt/nvme0n1/vime_ci/models:/root/models \ - -v /mnt/nvme0n1/vime_ci/datasets:/root/datasets \ - -w "$GITHUB_WORKSPACE" \ - inferactinc/public:vime-latest \ - bash -lc ' - set -euo pipefail - pip install -e . --no-deps --break-system-packages - TEST_PATH="$TEST_FILE" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [[ -n "$TEST_ARGS" ]]; then - read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") - else - TEST_ARGS_ARRAY=() - fi - if [ "$NUM_GPUS" = "0" ]; then - python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - else - python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" - fi - ' diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml deleted file mode 100644 index 9c6de3c5b..000000000 --- a/.github/workflows/pre-commit.yml +++ /dev/null @@ -1,41 +0,0 @@ -name: pre-commit - -on: - push: - branches: [main] - pull_request: - types: [opened, synchronize, reopened, ready_for_review] - -permissions: - contents: read - -jobs: - run-pre-commit: - name: Run pre-commit - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: '3.10' - cache: 'pip' - - - name: Install pre-commit - run: pip install --upgrade pip pre-commit - - - name: Cache pre-commit environments - uses: actions/cache@v5 - with: - path: ~/.cache/pre-commit - key: pre-commit-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }} - restore-keys: | - pre-commit-${{ runner.os }}- - - - name: Run pre-commit on all files - run: pre-commit run --all-files --show-diff-on-failure --color=always - diff --git a/.gitignore b/.gitignore index e9a268d45..c270e3a47 100644 --- a/.gitignore +++ b/.gitignore @@ -180,6 +180,7 @@ settings.json wandb/ outputs/ local/ +train_rollout_diff_exp/results/ **/rollout_data/ **/buffer_stats/ *.out diff --git a/README.md b/README.md index b400b4f9b..275531f0f 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,7 @@ [中文版](./README_zh.md) · [Repository](https://github.com/vllm-project/vime) +[![Documentation](https://img.shields.io/badge/docs-latest-brightgreen.svg?style=flat)](https://docs.vllm.ai/projects/vime/en/latest/) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/vllm-project/vime) **Vime** is an LLM post-training framework for RL scaling, built on [slime](https://github.com/THUDM/slime). It keeps slime's training stack and data-generation design while using [**vLLM**](https://github.com/vllm-project/vllm) (with [vllm-router](https://github.com/vllm-project/router)) as the default rollout backend. Vime provides two core capabilities: @@ -31,7 +32,9 @@ The vLLM community horizontally supports many LLM post-training frameworks, incl - [Table of Contents](#table-of-contents) - [Architecture Overview](#architecture-overview) - [Quick Start](#quick-start) + - [Agentic RL examples](#agentic-rl-examples) - [Arguments Walkthrough](#arguments-walkthrough) + - [Code Reading Path](#code-reading-path) - [Developer Guide](#developer-guide) - [slime doc](#slime-doc) - [FAQ](#faq) @@ -45,8 +48,8 @@ The vLLM community horizontally supports many LLM post-training frameworks, incl **Module Descriptions**: - **training (Megatron)**: Responsible for the main training process, reads data from the Data Buffer, and synchronizes parameters to the rollout module after training. -- **rollout (vLLM + router)**: Launches vLLM inference engines and routes generation requests; produces new data (including rewards/verifier outputs) and stores it in the Data Buffer. -- **data buffer**: A bridge module that manages prompt initialization, custom data, and rollout generation methods. +- **rollout (vLLM + router)**: Launches vLLM inference engines and routes generation requests; custom generate functions can wrap generation with multi-turn loops, tool calls, environment/sandbox interaction, and verifier-based rewards. +- **data buffer**: A bridge module that manages prompt initialization, custom data, and rollout generation methods, including agentic workflows that produce samples through the same interface. ## Quick Start @@ -56,6 +59,16 @@ For a comprehensive quick start guide covering environment setup, data preparati We also provide examples for some use cases not covered in the quick start guide; please check [examples](examples/). +### Agentic RL examples + +Agentic workloads use the standard rollout / Data Buffer loop through Vime's customization interfaces; they are not a separate framework: + +- [`examples/multi_agent`](examples/multi_agent/README.md): Multi-agent generation through `--custom-generate-function-path`. +- [`examples/fully_async`](examples/fully_async/README.md): Fully asynchronous rollout for long-tail agent generation. +- [`examples/coding_agent_rl`](examples/coding_agent_rl/README.md): End-to-end coding-agent RL with Claude Code or Codex, sandboxed tool use, test-based rewards, and token-correct trajectory segments. + +See the [Agentic RL Training Roadmap](docs/en/get_started/agent.md) and [Customization Guide](docs/en/get_started/customization.md). The coding-agent example ships an E2B-compatible backend, while the shared `vime.agent.sandbox.Sandbox` contract can be implemented for Docker, Modal, or local VMs. + ## Arguments Walkthrough Arguments in Vime are divided into three categories: @@ -68,6 +81,23 @@ Arguments in Vime are divided into three categories: For complete usage instructions, please refer to the [Usage Documentation](docs/en/get_started/usage.md). +## Code Reading Path + +Start from the training loop and follow the calls only as deep as needed: + +```text +train.py: train +├─ vime/ray/placement_group.py Ray resource and worker initialization +├─ vime/ray/rollout.py RolloutManager.generate: rollout orchestration +│ └─ vime/rollout/vllm_rollout.py Sample generation and reward computation +└─ vime/ray/actor_group.py RayTrainGroup.async_train: training dispatch + └─ vime/backends/megatron_utils/actor.py + ├─ model.py Megatron model execution + └─ loss.py RL losses and advantages +``` + +On a first pass, treat `vime/utils/arguments.py` as the configuration entry point. The deployment details in `vime/backends/vllm_utils/` and the weight-sync implementations under `vime/backends/megatron_utils/update_weight/` can wait until you need to change those areas. + ## Developer Guide - **Contributions are welcome!** If you have suggestions for new features, performance tuning, or feedback on user experience, feel free to submit an Issue or PR. diff --git a/README_zh.md b/README_zh.md index f2f6721fc..3b7a55b4d 100644 --- a/README_zh.md +++ b/README_zh.md @@ -2,6 +2,7 @@ [English](./README.md) · [代码仓库](https://github.com/vllm-project/vime) +[![文档](https://img.shields.io/badge/docs-latest-brightgreen.svg?style=flat)](https://docs.vllm.ai/projects/vime/zh-cn/latest/) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/vllm-project/vime) **Vime** 是基于 [slime](https://github.com/THUDM/slime) 的 RL scaling 用 LLM post-training 框架。在保留 slime 训练栈与数据生成设计的同时，默认以 [**vLLM**](https://github.com/vllm-project/vllm)（配合 [vllm-router](https://github.com/vllm-project/router)）作为 rollout 后端。Vime 提供两大核心能力： @@ -31,7 +32,9 @@ vLLM 社区横向支持许多 LLM post-training 框架，包括（按字母顺 - [目录](#目录) - [架构总览](#架构总览) - [快速开始](#快速开始) + - [Agentic RL 示例](#agentic-rl-示例) - [参数说明](#参数说明) + - [代码阅读路径](#代码阅读路径) - [开发指南](#开发指南) - [slime doc](#slime-doc) - [FAQ](#faq) @@ -45,8 +48,8 @@ vLLM 社区横向支持许多 LLM post-training 框架，包括（按字母顺 **模块说明**： - **training (Megatron)**：负责主训练流程，从 Data Buffer 读取数据，训练完后将参数同步至 rollout 模块； -- **rollout (vLLM + router)**：启动 vLLM 推理引擎并路由生成请求，产出新数据（含 reward/verifier），存储至 Data Buffer； -- **data buffer**：桥梁模块，管理 prompt 初始化、自定义数据与 rollout 生成方法。 +- **rollout (vLLM + router)**：启动 vLLM 推理引擎并路由生成请求；自定义生成函数可以在其上封装多轮循环、工具调用、环境/沙盒交互和基于 verifier 的奖励； +- **data buffer**：桥梁模块，管理 prompt 初始化、自定义数据与 rollout 生成方法，包括通过同一接口产出样本的 agent 工作流。 ## 快速开始 @@ -56,18 +59,45 @@ vLLM 社区横向支持许多 LLM post-training 框架，包括（按字母顺 我们还提供了一些未在快速开始中覆盖的使用示例，请查看 [examples](examples/)。 +### Agentic RL 示例 + +Agent 工作负载通过 Vime 的定制接口接入标准 rollout / Data Buffer 循环，并不是独立框架： + +- [`examples/multi_agent`](examples/multi_agent/README.md)：通过 `--custom-generate-function-path` 实现多 agent 生成； +- [`examples/fully_async`](examples/fully_async/README.md)：面向长尾 agent 生成的全异步 rollout； +- [`examples/coding_agent_rl`](examples/coding_agent_rl/README.md)：使用 Claude Code 或 Codex、沙盒工具、测试奖励和 token 精确轨迹片段的端到端 coding-agent RL； + +请参阅 [Agentic RL 训练路线图](docs/zh/get_started/agent.md)和[定制化指南](docs/zh/get_started/customization.md)。Coding-agent 示例内置 E2B 兼容后端；共享的 `vime.agent.sandbox.Sandbox` 协议也可以由 Docker、Modal 或本地虚拟机实现。 + ## 参数说明 Vime 的参数分为三类： 1. **Megatron 参数**：Vime 会读取 Megatron 中的全部参数，可通过传入如 `--tensor-model-parallel-size 2` 的方式配置 Megatron； -2. **vLLM 参数**：vLLM server 与 engine 相关选项以 `--vllm-` 为前缀（例如 `--vllm-gpu-memory-utilization`）。路由相关选项分两类前缀：vllm-router 自身的选项以 `--router-` 传入（例如 `--router-policy round_robin`、`--router-request-timeout-secs`），Vime 侧用于告诉 Vime *router 在哪里* 的编排参数则以 `--vllm-router-` 为前缀（`--vllm-router-ip`、`--vllm-router-port`）。完整参数见 [vime/backends/vllm_utils/arguments.py](vime/backends/vllm_utils/arguments.py)。 +2. **vLLM 参数**：vLLM server 与 engine 相关选项以 `--vllm-` 为前缀（例如 `--vllm-gpu-memory-utilization`）。vllm-router 自身的选项以 `--router-` 传入（例如 `--router-policy round_robin`）；Vime 侧的 router 编排参数使用 `--vllm-router-` 前缀，包括 `--vllm-router-ip`、`--vllm-router-port` 和实际生效的 `--vllm-router-request-timeout-secs`。完整参数见 [vime/backends/vllm_utils/arguments.py](vime/backends/vllm_utils/arguments.py)。 3. **框架参数**：与 Vime 编排相关的开关（rollout GPU、数据路径、RL 算法等），见 [vime/utils/arguments.py](vime/utils/arguments.py)。 `--rollout-num-gpus-per-engine` 对应每个 vLLM engine 的 tensor parallel size。默认 rollout 入口为 `vime.rollout.vllm_rollout.generate_rollout`。 完整使用说明请查阅 [使用文档](docs/zh/get_started/usage.md)。 +## 代码阅读路径 + +建议从训练主循环开始，只在需要时继续深入： + +```text +train.py: train +├─ vime/ray/placement_group.py Ray 资源与 worker 初始化 +├─ vime/ray/rollout.py RolloutManager.generate：rollout 编排 +│ └─ vime/rollout/vllm_rollout.py 样本生成与奖励计算 +└─ vime/ray/actor_group.py RayTrainGroup.async_train：训练调度 + └─ vime/backends/megatron_utils/actor.py + ├─ model.py Megatron 模型执行 + └─ loss.py RL loss 与 advantage +``` + +第一次阅读时，可以把 `vime/utils/arguments.py` 当作配置入口。只有需要修改相关区域时，再深入 `vime/backends/vllm_utils/` 的部署细节和 `vime/backends/megatron_utils/update_weight/` 下的权重同步实现。 + ## 开发指南 - **欢迎贡献！** 若有功能建议、性能调优或使用体验反馈，欢迎提交 Issue / PR。 diff --git a/docker/Dockerfile b/docker/Dockerfile index aca34ae6e..a6acb47b9 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,13 +1,20 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0-cu129-ubuntu2404 +ARG BASE_IMAGE=vllm/vllm-openai:nightly FROM ${BASE_IMAGE} # ======================================== Arguments ============================================= ARG PATCH_VERSION=latest ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 +ARG DEEPGEMM_COMMIT=b38a77cd193cf38f670caae192310521d24343be +ARG DEEPEP_COMMIT=6845ffd9d59126ec0030c13e0e155935a61e5b5a +ARG DEEPEP_CUDA_ARCH_LIST=9.0;10.0;10.3 +ARG FLASH_QLA_COMMIT=821fd9d37ede18fdc2a4e707fefe3770bfc32e58 +ARG TRANSFORMER_ENGINE_COMMIT=c9877beb87ad7e711e1869dd0b5062167ede447a +ARG TRANSFORMER_ENGINE_CUDA_ARCHS=90;100a;103a +ARG TMS_COMMIT=8d30c59ca12a68d9deccbc9c6599076a1218cbc5 -ARG ENABLE_CUDA_13=0 -ARG TMS_CUDA_MAJOR= +ARG ENABLE_CUDA_13=1 +ARG FA2_MAX_JOBS=64 # ======================================== Setup ============================================= @@ -15,13 +22,22 @@ WORKDIR /root/ # ======================================== Apt dependencies ============================================= -# vllm/vllm-openai base is an inference image — add cu12 dev headers + cmake/git +# vllm/vllm-openai base is an inference image — add CUDA dev headers + cmake/git # so TE / apex / flash-attn source builds find cusparse.h etc. +# Dev packages must match the base toolkit: -12-9 for cu129, -13-0 for cu130. RUN apt-get update && apt-get install -y \ - nvtop rsync dnsutils git cmake \ - cuda-nvrtc-dev-12-9 cuda-nvml-dev-12-9 cuda-profiler-api-12-9 cuda-nvtx-12-9 \ - libcusparse-dev-12-9 libcusolver-dev-12-9 libcufft-dev-12-9 libcurand-dev-12-9 \ - libcudnn9-dev-cuda-12 && \ + nvtop rsync dnsutils prometheus git cmake && \ + if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + apt-get install -y \ + cuda-nvrtc-dev-13-0 cuda-nvml-dev-13-0 cuda-profiler-api-13-0 cuda-nvtx-13-0 \ + libcusparse-dev-13-0 libcusolver-dev-13-0 libcufft-dev-13-0 libcurand-dev-13-0 \ + libcudnn9-dev-cuda-13; \ + else \ + apt-get install -y \ + cuda-nvrtc-dev-12-9 cuda-nvml-dev-12-9 cuda-profiler-api-12-9 cuda-nvtx-12-9 \ + libcusparse-dev-12-9 libcusolver-dev-12-9 libcufft-dev-12-9 libcurand-dev-12-9 \ + libcudnn9-dev-cuda-12; \ + fi && \ rm -rf /var/lib/apt/lists/* # vllm/vllm-openai base only ships python3; subsequent source builds invoke `python`. @@ -29,56 +45,46 @@ RUN ln -sf /usr/bin/python3 /usr/local/bin/python # ====================================== Python dependencies ============================================ -# The compilation is slow, thus should be put at top -# TransformerEngines does not support too high FA2 -RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation +# The validated TransformerEngine 2.16 context-parallel stack uses FA2 + FA3. +RUN pip uninstall -y flash-attn-4 flash_attn_4 || true +RUN MAX_JOBS=${FA2_MAX_JOBS} pip -v install flash-attn==2.8.3 --no-build-isolation -# The compilation is slow, thus should be put at top +# This FA3 commit provides the window_size_left/window_size_right API used by TE 2.16. RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ - cd flash-attention/ && git checkout fbf24f67cf7f6442c5cfb2c1057f4bfc57e72d89 && git submodule update --init && cd hopper/ && \ - MAX_JOBS=96 python setup.py install && \ - export python_path=`python -c "import site; print(site.getsitepackages()[0])"` && \ - mkdir -p $python_path/flash_attn_3 && \ - cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py && \ - rm -rf flash-attention/ - -RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps + cd flash-attention/ && git checkout 002cce0a1068f8c07dfccb5a1d232b9a3276947c && git submodule update --init && \ + cd hopper/ && \ + FLASH_ATTENTION_FORCE_BUILD=TRUE MAX_JOBS=96 pip -v install . --no-build-isolation && \ + cd /root/ && rm -rf flash-attention/ -RUN pip install flash-linear-attention==0.4.1 +RUN pip install flash-linear-attention==0.4.2 +# FlashQLA: optional GDN backend for Qwen3.5/Qwen3-Next (--qwen-gdn-backend flashqla; requires SM90+) ARG INSTALL_FLASHQLA=0 RUN if [ "${INSTALL_FLASHQLA}" = "1" ]; then \ - pip install git+https://github.com/QwenLM/FlashQLA.git --no-build-isolation; \ + pip install git+https://github.com/QwenLM/FlashQLA.git@${FLASH_QLA_COMMIT} --no-build-isolation; \ else \ echo "Skipping FlashQLA (INSTALL_FLASHQLA=0; sm90/Hopper-only — use --qwen-gdn-backend fla)"; \ fi -RUN pip install tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ - -# cublas + (cu13) cuda dev headers. The arm64 (sbsa) vllm/vllm-openai base ships -# the cublas runtime .so but not the dev header (cublas_v2.h) / unversioned .so -# symlink that the x86 base has; TE needs the header and its CMake only creates -# the CUDA::cublas target when both exist. The dev pkg must match the toolkit -# CUDA major. For cu13 the top apt block's -12-9 dev set is the wrong major, so -# also install the -13-0 nvrtc/nvtx/etc. headers TE's nvcc build needs. Placed -# after the slow flash-attn layers so their cache is preserved; before TE. +# cublas dev header for TE CMake (arm64 base ships runtime .so but not the header). +# cu13 also needs the -13-0 headers that TE's nvcc build expects. RUN apt-get update && \ if [ "${ENABLE_CUDA_13}" = "1" ]; then \ - apt-get install -y libcublas-dev-13-0 \ - cuda-nvrtc-dev-13-0 cuda-nvtx-13-0 cuda-nvml-dev-13-0 cuda-profiler-api-13-0 \ - libcusparse-dev-13-0 libcusolver-dev-13-0 libcufft-dev-13-0 libcurand-dev-13-0; \ - else apt-get install -y libcublas-dev-12-9; fi && \ + apt-get install -y libcublas-dev-13-0; \ + else \ + apt-get install -y libcublas-dev-12-9; \ + fi && \ rm -rf /var/lib/apt/lists/* -# TE does not have wheel on cuda 13 yet, thus need to install from source. -# TE is built with --no-build-isolation, so its build deps must be pre-installed; -# install the set TE's release_v2.10 CI uses (wheel packaging ninja pybind11). -# nvidia-mathdx is left unpinned (as in TE CI): 26.6.0 is x86-only, arm64 (sbsa) -# max is 25.6.0, and TE release_v2.10 does not pin or reference it (it links plain -# CUDA::cublas), so the resolver picking the per-arch latest is safe. +# The cu13 TE wheel is built against a newer cuBLAS than the vLLM CUDA 13.0 +# base. Build the pinned source revision against this image. RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ - pip install nvidia-mathdx pybind11 ninja wheel packaging && \ - pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ + pip install pybind11 && \ + NCCL_INCLUDE_DIR="$(python -c 'import importlib.util; print(next(iter(importlib.util.find_spec("nvidia.nccl").submodule_search_locations)) + "/include")')" && \ + export CPATH="${NCCL_INCLUDE_DIR}${CPATH:+:${CPATH}}" && \ + NVTE_CUDA_ARCHS="${TRANSFORMER_ENGINE_CUDA_ARCHS}" MAX_JOBS=64 \ + pip -v install --no-build-isolation \ + git+https://github.com/NVIDIA/TransformerEngine.git@${TRANSFORMER_ENGINE_COMMIT}; \ else \ - pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.16.1"; \ fi RUN NVCC_APPEND_FLAGS="--threads 4" \ @@ -89,43 +95,49 @@ RUN NVCC_APPEND_FLAGS="--threads 4" \ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ cd Megatron-LM && git checkout ${MEGATRON_COMMIT} -# torch_memory_saver pinned to a193d9dd (upstream slime #1916). The newer commit -# ships a multi-CUDA wheel and requires TMS_CUDA_MAJOR at build time; default it -# to the running torch's CUDA major (slime #1924). -RUN TMS_CUDA_MAJOR="${TMS_CUDA_MAJOR:-$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')}" && \ +# zhuzilin fork builds, grouped together right after Megatron-LM: +# torch_memory_saver, plus the GLM-5 train/rollout alignment kernels. +RUN TMS_CUDA_MAJOR="$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')" && \ export TMS_CUDA_MAJOR && \ - pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@a193d9dd1b877d33c64a41cfb3db9f867df2d926 --no-cache-dir --force-reinstall -RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation + pip install git+https://github.com/zhuzilin/torch_memory_saver.git@${TMS_COMMIT} --no-cache-dir --force-reinstall + +RUN git clone https://github.com/zhuzilin/DeepGEMM.git --recursive && \ + cd DeepGEMM && git checkout ${DEEPGEMM_COMMIT} && \ + git submodule update --init --recursive && \ + bash build_sgl_deep_gemm.sh && \ + pip install --force-reinstall --no-deps dist/sgl_deep_gemm-*.whl && \ + cd /root/ && rm -rf DeepGEMM + +RUN git clone https://github.com/zhuzilin/DeepEP.git /root/DeepEP && \ + cd /root/DeepEP && git checkout ${DEEPEP_COMMIT} && \ + CUDA_ARCHS="${DEEPEP_CUDA_ARCH_LIST}" && \ + if [ -d /usr/local/cuda/include/cccl ]; then \ + export CPATH="/usr/local/cuda/include/cccl${CPATH:+:$CPATH}"; \ + export NVCC_PREPEND_FLAGS="-I/usr/local/cuda/include/cccl ${NVCC_PREPEND_FLAGS:-}"; \ + fi && \ + TORCH_CUDA_ARCH_LIST="${CUDA_ARCHS}" MAX_JOBS=64 python setup.py bdist_wheel && \ + pip install --force-reinstall --no-deps dist/deep_ep-*.whl && \ + cd /root/ && rm -rf DeepEP RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation -# This patch from masahi will be included in later Triton releases -RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ - (cd /root && git clone -b feat/v350_plus_8045 https://github.com/fzyzcjy/triton.git && cd triton && pip install -r python/requirements.txt && pip install --verbose -e .); \ - fi - COPY requirements.txt /tmp/requirements.txt RUN pip install --ignore-installed PyJWT && \ pip install -r /tmp/requirements.txt # https://github.com/pytorch/pytorch/issues/168167 -RUN pip install nvidia-cudnn-cu12==9.16.0.29 - -# reinstall numpy 1.x for megatron -RUN pip install "numpy<2" - -RUN pip install IPython +# cu130 base already ships nvidia-cudnn-cu13 9.19.0; only cu129 needs the pin. +RUN if [ "${ENABLE_CUDA_13}" != "1" ]; then \ + pip install nvidia-cudnn-cu12==9.16.0.29; \ + fi -# Pin vllm-router explicitly so the vllm rollout routing layer is a visible build step -# (also in requirements.txt; pulling it here makes the layer cache-able and fail-fast). -RUN pip install "vllm-router>=0.1.14" +# Megatron requires NumPy 1.x, while SciPy 1.18+ requires NumPy 2.x. +RUN pip install "numpy==1.26.4" "scipy==1.17.1" -RUN rm -rf /root/.cache/pip +RUN rm -rf /root/.cache/pip /root/flash-attention # ====================================== Patches ============================================ -# Patch megatron BEFORE pip install -e . so any patch hunks that touch setup.py -# or C++/CUDA extensions are picked up by the build. -COPY docker/patch/${PATCH_VERSION}/megatron.patch /root/Megatron-LM/ +COPY docker/patch/${PATCH_VERSION}/megatron*.patch /root/Megatron-LM/ RUN cd Megatron-LM && \ git update-index --refresh && \ git apply megatron.patch --3way && \ @@ -133,24 +145,25 @@ RUN cd Megatron-LM && \ echo "Patch failed to apply cleanly. Please resolve conflicts." && \ exit 1; \ fi && \ - rm megatron.patch && \ + rm -f megatron*.patch && \ pip install -e . -# Patch vLLM: skip execute_dummy_batch while the engine is asleep / being put to sleep, so a -# colocate DP+EP rollout's staged sleep/wake weight-sync doesn't run a forward against freed KV -# (illegal memory access at update_weights / the post-rollout sleep offload). vLLM PR #44483 -# lineage, broadened to guard on EngineCore.is_sleeping() (covers the cuMem offload window). -# vLLM is a pip install (not a git checkout) so apply with plain `git apply` (no --3way). +# Patch vLLM with vime's local fixes. vLLM is a pip install (not a git checkout) +# so apply with plain `git apply` (no --3way). Pull-weights lands first because +# the general patch also updates gpu_worker.py against the resulting line layout. +COPY docker/patch/${PATCH_VERSION}/vllm-pull_weights.patch /tmp/vllm-pull_weights.patch COPY docker/patch/${PATCH_VERSION}/vllm.patch /tmp/vllm.patch RUN VLLM_SITE="$(python3 -c 'import os, vllm; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" && \ cd "$VLLM_SITE" && \ - git apply -v /tmp/vllm.patch && \ - rm /tmp/vllm.patch + git apply -v /tmp/vllm-pull_weights.patch && \ + git apply -v --allow-empty /tmp/vllm.patch && \ + rm /tmp/vllm-pull_weights.patch /tmp/vllm.patch # ====================================== Install main package ============================================ +ARG VIME_REPO=https://github.com/vllm-project/vime.git ARG VIME_COMMIT=main -RUN git clone https://github.com/vllm-project/vime.git /root/vime && \ +RUN git clone ${VIME_REPO} /root/vime && \ cd /root/vime && \ git checkout ${VIME_COMMIT} && \ pip install -e . --no-deps @@ -158,21 +171,8 @@ RUN git clone https://github.com/vllm-project/vime.git /root/vime && \ RUN cd /root/vime/vime/backends/megatron_utils/kernels/int4_qat && \ pip install . --no-build-isolation -# ====================================== Build-time smoke ============================================ - -# Fail-fast import smoke + flashinfer version pin check. Catches ABI / version -# regressions at build time instead of first GPU run. -RUN python3 -c "\ -import vllm, vime, flashinfer; \ -from vime.backends.vllm_utils.vllm_engine import VLLMEngine; \ -print('vllm', vllm.__version__); \ -print('flashinfer', flashinfer.__version__); \ -print('VLLMEngine import ok')" - -# Megatron stream serialization. vllm rollout + megatron actor co-located on -# the same GPU need this; default value causes stream-level contention. -ENV CUDA_DEVICE_MAX_CONNECTIONS=1 - -# Reset ENTRYPOINT inherited from vllm/vllm-openai base (`vllm serve`). +# Reset ENTRYPOINT inherited from the vllm/vllm-openai base (`vllm serve`), so the +# image is a plain bash/ray environment. Without this, `docker run ... bash -c ...` +# and `ray job submit` append to `vllm serve` and break. ENTRYPOINT [] CMD ["/bin/bash"] diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 7ef525cde..f9e42c9aa 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,14 +1,14 @@ # syntax=docker/dockerfile:1.7 -ARG BASE_IMAGE=quay.io/ascend/vllm-ascend -ARG BASE_IMAGE_TAG=v0.22.1rc1-a3 +ARG BASE_IMAGE=quay.io/atlas-ci/vllm-ascend +ARG BASE_IMAGE_TAG=v0.28.0-fd81546-a3 FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} SHELL ["/bin/bash", "-o", "pipefail", "-c"] WORKDIR /root ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 -ARG MEGATRON_BRIDGE_COMMIT=7f0fb3456f8ffe47599b5fd167b454605d85f932 +ARG MEGATRON_BRIDGE_COMMIT=3fd3768045422d0aa5c97e90a4e6c659aea9acb9 ARG MINDSPEED_COMMIT=fc63de5c48426dd019c3b3f39e65f5bdf56e4086 ARG MEGATRON_ADAPTOR_COMMIT=15582addff3f3d4680e350826fa70d012b475509 ARG TRANSFORMER_ENGINE_NPU_COMMIT=d743c83d060d5edc48867ecb9e93ec80d81860e4 @@ -35,7 +35,7 @@ ENV SOC_VERSION=$SOC_VERSION \ # PATCH MAINTENANCE: keep patch COPY/apply operations and # docker/npu_patch/series.conf synchronized. COPY docker/npu_patch /opt/npu_patch -COPY docker/patch/latest/megatron.patch /opt/vime_patch/megatron.patch +COPY docker/patch/latest/megatron.patch /opt/npu_patch/megatron-common.patch RUN git config --global http.sslVerify false @@ -90,9 +90,9 @@ RUN git clone https://github.com/ISEEKYAN/mbridge.git /root/mbridge && \ # The NPU Megatron patch is based on the common Vime Megatron patch, so the # common patch must be applied first. RUN git -C /root/Megatron-LM apply --check --whitespace=nowarn \ - /opt/vime_patch/megatron.patch && \ + /opt/npu_patch/megatron-common.patch && \ git -C /root/Megatron-LM apply --whitespace=nowarn \ - /opt/vime_patch/megatron.patch && \ + /opt/npu_patch/megatron-common.patch && \ git -C /root/Megatron-LM apply --check --whitespace=nowarn \ /opt/npu_patch/megatron.patch && \ git -C /root/Megatron-LM apply --whitespace=nowarn \ @@ -132,39 +132,17 @@ RUN pip install \ # Vime currently imports torch_memory_saver from its training actor. RUN git clone --depth 1 --branch 2026.6.0 \ https://github.com/sgl-project/sgl-kernel-npu.git /root/sgl-kernel-npu && \ - cd /root/sgl-kernel-npu && \ - bash build.sh -a kernels && \ - bash build.sh -a memory-saver && \ - pip install --no-deps \ - output/torch_memory_saver-*.whl && \ + cd /root/sgl-kernel-npu/contrib/torch_memory_saver/python && \ + python3 setup.py bdist_wheel && \ + python3 -m pip install --no-deps \ + dist/torch_memory_saver-*.whl && \ cd /root && \ rm -rf /root/sgl-kernel-npu -# ---- fla-npu: NPU GDN (gated delta net) kernels for Qwen3.5-35B-A3B ------------ -RUN git clone -b v26.6.0 https://github.com/flashserve/flash-linear-attention-npu \ - /root/flash-linear-attention-npu && \ - git -C /root/flash-linear-attention-npu checkout 14c2c92 - -RUN pip uninstall -y torch_npu && \ - pip install torch_npu==2.10.0.post2 && \ - cd /root/flash-linear-attention-npu && \ - source /usr/local/Ascend/cann/set_env.sh && \ - printf 'y' | bash install_deps.sh && \ - pip install -r requirements.txt && \ - python3 scripts/check_npu_env.py --build-only - -RUN cd /root/flash-linear-attention-npu && \ - source /usr/local/Ascend/cann/set_env.sh && \ - bash build.sh --soc=ascend910_93 --pkg --vendor_name=fla_npu && \ - bash build_out/fla-npu-fla_npu_linux-aarch64.run && \ - export LD_LIBRARY_PATH=/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer/op_api/lib:${LD_LIBRARY_PATH} && \ - cd torch_custom/fla_npu && \ - bash build.sh - # Minimal import check. RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - python3 -c 'import megatron, mindspeed, megatron_adaptor, transformer_engine, torch_memory_saver, vime, vllm, vllm_ascend, fla_npu;' + python3 -c 'import megatron, mindspeed, megatron_adaptor, transformer_engine, torch_memory_saver, vime, vllm, vllm_ascend;' WORKDIR /root/vime ENTRYPOINT [] -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/docker/README.md b/docker/README.md index 21637d34c..45186da2b 100644 --- a/docker/README.md +++ b/docker/README.md @@ -1,7 +1,7 @@ # Docker release rule vime ships one image based on the official vllm image, published as -`inferactinc/public:vime-latest`. Supports GB200/300 and H100/200. +`vllm/vime:latest`. Supports GB200/300 and H100/200. Build locally: diff --git a/docker/justfile b/docker/justfile index 44f060614..ea786d20d 100644 --- a/docker/justfile +++ b/docker/justfile @@ -1,35 +1,39 @@ -# Release images publish to the company hub `inferactinc/public` as multi-arch +# Release images publish to `vllm/vime` on DockerHub as multi-arch # manifests, exactly like upstream `vllm/vllm-openai:nightly`: ONE tag carries # both the linux/amd64 and linux/arm64 digests and `docker pull` auto-selects # the platform. No arch suffix is ever published as a tag. # -# CUDA 12.9 is the default, so it carries NO cu marker in the tag. Only the -# non-default cu13 variant is suffixed. -# # The CUDA + source-built FA/TE images can't be cross-emulated, so each arch is # built on its own native host and pushed BY DIGEST (no tag lands in the hub), # then the two digests are fused into the final tag with `just manifest`. # +# CUDA 13 is the default and carries no cu marker. The explicit cu13 tags are +# compatibility aliases for users of the previous two-variant release scheme. +# # Tag scheme: -# inferactinc/public:vime- immutable, multi-arch (cu12.9) -# inferactinc/public:vime-latest rolling, multi-arch (cu12.9) -# inferactinc/public:vime-cu13- immutable, multi-arch (cu13 variant) -# inferactinc/public:vime-cu13-latest rolling, multi-arch (cu13 variant) +# vllm/vime: immutable, multi-arch (cu13) +# vllm/vime:latest rolling, multi-arch (cu13) +# vllm/vime:cu13- compatibility alias +# vllm/vime:cu13-latest compatibility alias # comes from docker/version.txt. -IMAGE := "inferactinc/public" -BUILDER := "vime-builder" +IMAGE := "vllm/vime" +BUILDER := env("VIME_BUILDER", "vime-builder") +VIME_COMMIT := env("VIME_COMMIT", "main") +FA2_MAX_JOBS := env("VIME_FA2_MAX_JOBS", "64") +DEEPEP_CUDA_ARCH_LIST := env("VIME_DEEPEP_CUDA_ARCH_LIST", "9.0;10.0;10.3") +TRANSFORMER_ENGINE_CUDA_ARCHS := env("VIME_TRANSFORMER_ENGINE_CUDA_ARCHS", "90;100a;103a") # ---- per-arch build, pushed BY DIGEST (run once on an amd64 host, once on an arm64 host) ---- -# Default — cu12.9, no cu marker in the tag. +# Default — latest vLLM nightly / CUDA 13 base. build: - ARG_TAG_SUFFIX="" ARG_BUILD_EXTRA_ARGS="--build-arg INSTALL_FLASHQLA=1" just _build-digest + ARG_TAG_SUFFIX="" ARG_BUILD_EXTRA_ARGS="--build-arg VIME_COMMIT={{ VIME_COMMIT }} --build-arg FA2_MAX_JOBS={{ FA2_MAX_JOBS }} --build-arg DEEPEP_CUDA_ARCH_LIST={{ DEEPEP_CUDA_ARCH_LIST }} --build-arg TRANSFORMER_ENGINE_CUDA_ARCHS={{ TRANSFORMER_ENGINE_CUDA_ARCHS }}" just _build-digest -# cu13 variant — vLLM default-CUDA base; ENABLE_CUDA_13 builds the CUDA-13 -# TransformerEngine/Triton on top. +# Compatibility target for the explicit cu13 aliases. Use the same nightly base +# as the default image so the aliases cannot drift from it. build-cu13: - ARG_TAG_SUFFIX="-cu13" ARG_BUILD_EXTRA_ARGS='--build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.0-ubuntu2404 --build-arg ENABLE_CUDA_13=1' just _build-digest + ARG_TAG_SUFFIX="-cu13" ARG_BUILD_EXTRA_ARGS='--build-arg BASE_IMAGE=vllm/vllm-openai:nightly --build-arg ENABLE_CUDA_13=1 --build-arg VIME_COMMIT={{ VIME_COMMIT }} --build-arg FA2_MAX_JOBS={{ FA2_MAX_JOBS }} --build-arg DEEPEP_CUDA_ARCH_LIST={{ DEEPEP_CUDA_ARCH_LIST }} --build-arg TRANSFORMER_ENGINE_CUDA_ARCHS={{ TRANSFORMER_ENGINE_CUDA_ARCHS }}' just _build-digest _build-digest: #!/bin/bash @@ -40,39 +44,38 @@ _build-digest: META="/tmp/vime${ARG_TAG_SUFFIX}-$(uname -m).json" # push-by-digest requires the docker-container buildx driver. - docker buildx inspect {{BUILDER}} >/dev/null 2>&1 || docker buildx create --name {{BUILDER}} --driver docker-container + docker buildx inspect {{ BUILDER }} >/dev/null 2>&1 || docker buildx create --name {{ BUILDER }} --driver docker-container - docker buildx build --builder {{BUILDER}} --platform "$PLATFORM" -f docker/Dockerfile $ARG_BUILD_EXTRA_ARGS --build-arg HTTP_PROXY="$http_proxy" --build-arg HTTPS_PROXY="$https_proxy" --build-arg NO_PROXY="localhost,127.0.0.1" -o "type=image,name={{IMAGE}},push-by-digest=true,push=true" --metadata-file "$META" . + docker buildx build --builder {{ BUILDER }} --platform "$PLATFORM" -f docker/Dockerfile $ARG_BUILD_EXTRA_ARGS --build-arg HTTP_PROXY="${http_proxy:-}" --build-arg HTTPS_PROXY="${https_proxy:-}" --build-arg NO_PROXY="localhost,127.0.0.1" -o "type=image,name={{ IMAGE }},push-by-digest=true,push=true" --metadata-file "$META" . echo "Pushed vime${ARG_TAG_SUFFIX} ${PLATFORM} by digest — pass this to \`just manifest\`:" jq -r '."containerimage.digest"' "$META" # ---- fuse the two per-arch digests into one multi-arch tag ---- -# Run once after `build` (or `build-cu13`) has pushed on BOTH hosts, passing the -# digests it printed. For the default cu12.9 image leave VARIANT empty: -# just manifest "" sha256: sha256: -> vime- + vime-latest -# just manifest cu13 sha256: sha256: -> vime-cu13- + vime-cu13-latest +# Run once after `build` (or `build-cu13`) has pushed on both hosts. Publish the +# default tags first; the same digests can also publish the cu13 aliases: +# just manifest "" sha256: sha256: -> vime: + vime:latest +# just manifest cu13 sha256: sha256: -> cu13 compatibility aliases manifest VARIANT AMD_DIGEST ARM_DIGEST: #!/bin/bash set -euxo pipefail cd .. VERSION="$(cat docker/version.txt | tr -d '\n')" - SUFFIX="" - [ -n "{{VARIANT}}" ] && SUFFIX="-{{VARIANT}}" - docker buildx imagetools create -t "{{IMAGE}}:vime${SUFFIX}-${VERSION}" -t "{{IMAGE}}:vime${SUFFIX}-latest" "{{IMAGE}}@{{AMD_DIGEST}}" "{{IMAGE}}@{{ARM_DIGEST}}" + PREFIX="" + [ -n "{{ VARIANT }}" ] && PREFIX="{{ VARIANT }}-" + docker buildx imagetools create -t "{{ IMAGE }}:${PREFIX}${VERSION}" -t "{{ IMAGE }}:${PREFIX}latest" "{{ IMAGE }}@{{ AMD_DIGEST }}" "{{ IMAGE }}@{{ ARM_DIGEST }}" -# ---- single-arch test/debug image for the run-ci-image validation job ---- -# The e2e-test-image runner is x86, so this is amd64-only and -# needs no manifest. +# ---- single-arch test/debug image for Buildkite GPU validation ---- +# This target is amd64-only and needs no manifest. build-test: #!/bin/bash set -euxo pipefail cd .. VERSION="$(cat docker/version.txt | tr -d '\n')" - docker build -f docker/Dockerfile . --build-arg HTTP_PROXY="$http_proxy" --build-arg HTTPS_PROXY="$https_proxy" --build-arg NO_PROXY="localhost,127.0.0.1" --build-arg INSTALL_FLASHQLA=1 -t "{{IMAGE}}:vime-test-${VERSION}" - docker push "{{IMAGE}}:vime-test-${VERSION}" + docker build -f docker/Dockerfile . --build-arg HTTP_PROXY="${http_proxy:-}" --build-arg HTTPS_PROXY="${https_proxy:-}" --build-arg NO_PROXY="localhost,127.0.0.1" --build-arg DEEPEP_CUDA_ARCH_LIST={{ DEEPEP_CUDA_ARCH_LIST }} --build-arg TRANSFORMER_ENGINE_CUDA_ARCHS={{ TRANSFORMER_ENGINE_CUDA_ARCHS }} -t "{{ IMAGE }}:test-${VERSION}" + docker push "{{ IMAGE }}:test-${VERSION}" - docker tag "{{IMAGE}}:vime-test-${VERSION}" "{{IMAGE}}:vime-test-latest" - docker push "{{IMAGE}}:vime-test-latest" + docker tag "{{ IMAGE }}:test-${VERSION}" "{{ IMAGE }}:test-latest" + docker push "{{ IMAGE }}:test-latest" diff --git a/docker/npu_patch/README.md b/docker/npu_patch/README.md index d9656ae89..48b351166 100644 --- a/docker/npu_patch/README.md +++ b/docker/npu_patch/README.md @@ -1,135 +1,147 @@ # Vime NPU Patch Installation Guide -This guide provides instructions for installing Vime with NPU support, including all required dependencies and patches. +This guide provides instructions for installing Vime with NPU support, including the required dependencies and patches. ## Component Version Mapping -| Component | Version/Commit | Source | -| --------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------- | -| vime | main | [GitHub](https://github.com/vllm-project/vime/tree/main) | -| Megatron-Bridge | 7f0fb3456f8ffe47599b5fd167b454605d85f932 | [GitHub](https://github.com/radixark/Megatron-Bridge) | -| Megatron-LM | 1dcf0dafa884ad52ffb243625717a3471643e087 | [GitHub](https://github.com/NVIDIA/Megatron-LM) | -| MegatronAdaptor | 15582addff3f3d4680e350826fa70d012b475509 | [GitCode](https://gitcode.com/Ascend/MegatronAdaptor) | -| TransformerEngineNPU | d743c83d060d5edc48867ecb9e93ec80d81860e4 | [GitCode](https://gitcode.com/Ascend/TransformerEngineNPU) | -| MindSpeed | fc63de5c48426dd019c3b3f39e65f5bdf56e4086 | [GitCode](https://gitcode.com/Ascend/MindSpeed) | -| HDK | 25.3.RC1 | [Ascend](https://www.hiascend.com/hardware/firmware-drivers/commercial?product=7\&model=33) | -| CANN | 9.0.0 | [Ascend](https://www.hiascend.com/developer/download/community/result?module=cann\&cann=9.0.0\&product=7\&model=33) | +| Component | Version/Commit | Source | +| --- | --- | --- | +| Base image | `v0.28.0-fd81546-a3` | `quay.io/atlas-ci/vllm-ascend` | +| vLLM | `e6bfe03ad73a3330cb427885aa90d97a12e1c704` | [GitHub](https://github.com/vllm-project/vllm) | +| vLLM-Ascend | `fd815467c221ee600137f6bdd53fe354d5e7c999` | [GitHub](https://github.com/vllm-project/vllm-ascend) | +| Megatron-LM | `1dcf0dafa884ad52ffb243625717a3471643e087` | [GitHub](https://github.com/NVIDIA/Megatron-LM) | +| Megatron-Bridge | `3fd3768045422d0aa5c97e90a4e6c659aea9acb9` | [GitHub](https://github.com/radixark/Megatron-Bridge) | +| mbridge | `89eb10887887bc74853f89a4de258c0702932a1c` | [GitHub](https://github.com/ISEEKYAN/mbridge) | +| MegatronAdaptor | `15582addff3f3d4680e350826fa70d012b475509` | [GitCode](https://gitcode.com/Ascend/MegatronAdaptor) | +| TransformerEngineNPU | `d743c83d060d5edc48867ecb9e93ec80d81860e4` | [GitCode](https://gitcode.com/Ascend/TransformerEngineNPU) | +| MindSpeed | `fc63de5c48426dd019c3b3f39e65f5bdf56e4086` | [GitCode](https://gitcode.com/Ascend/MindSpeed) | +| torch_memory_saver (NPU) | `sgl-kernel-npu` tag `2026.6.0` | [GitHub](https://github.com/sgl-project/sgl-kernel-npu) | ## Preparing the Running Environment -Run the steps below in a Python 3.12 environment with CANN 9.0.0. A -`quay.io/ascend/vllm-ascend:nightly-main-a3` container can be used as the base. +Run the following steps inside the base image listed above, with the Ascend devices and host driver mounted. The base image provides Python, CANN, PyTorch, torch-npu, vLLM and vLLM-Ascend. + +Use a checkout of this Vime revision at `/root/vime`. Start with unpatched dependency source trees; do not repeat these steps in an already-patched Vime image. ```bash -export WORKSPACE=/root -cd "${WORKSPACE}" +export VIME_INSTALL_ROOT=/root +export PATCH_DIR="${VIME_INSTALL_ROOT}/vime/docker/npu_patch" +source /usr/local/Ascend/ascend-toolkit/set_env.sh ``` -Vime's Ascend NPU adaptation lives on the **`ascend`** branch, so clone that -branch (not `main`): +Preserve the base image's serving package versions when installing training dependencies: ```bash -git clone --branch ascend https://github.com/vllm-project/vime.git "${WORKSPACE}/vime" -export PATCH_DIR="${WORKSPACE}/vime/docker/npu_patch" +export PIP_CONSTRAINT="$(mktemp /tmp/vime-npu-constraints.XXXXXX)" +python3 - <<'PY' +import importlib.metadata as metadata +import os + +names = [ + "numpy", "ray", "torch", "torch-npu", "torchvision", + "transformers", "triton-ascend", "vllm", "vllm-ascend", +] +with open(os.environ["PIP_CONSTRAINT"], "w") as constraints: + constraints.write("\n".join(f"{name}=={metadata.version(name)}" for name in names) + "\n") +PY ``` -#### 1. Megatron-Bridge +### 1. vLLM and vLLM-Ascend -Used via `PYTHONPATH` (no editable install); it requires `nvidia-modelopt`. +Both packages are installed in editable mode in the base image. Apply the patches to their existing source trees: ```bash -export MEGATRON_BRIDGE_COMMIT=7f0fb3456f8ffe47599b5fd167b454605d85f932 -export MBRIDGE_COMMIT=89eb10887887bc74853f89a4de258c0702932a1c -pip install "git+https://github.com/ISEEKYAN/mbridge.git@${MBRIDGE_COMMIT}" --no-deps -git clone --branch bridge https://github.com/radixark/Megatron-Bridge.git "${WORKSPACE}/Megatron-Bridge" -git -C "${WORKSPACE}/Megatron-Bridge" checkout "${MEGATRON_BRIDGE_COMMIT}" - -git -C "${WORKSPACE}/Megatron-Bridge" apply --whitespace=nowarn "${PATCH_DIR}/megatron-bridge.patch" +git -C /vllm-workspace/vllm apply --check "${PATCH_DIR}/vllm.patch" +git -C /vllm-workspace/vllm apply "${PATCH_DIR}/vllm.patch" -pip install --no-build-isolation "nvidia-modelopt[torch]>=0.37.0" +git -C /vllm-workspace/vllm-ascend apply --check "${PATCH_DIR}/vllm-ascend.patch" +git -C /vllm-workspace/vllm-ascend apply "${PATCH_DIR}/vllm-ascend.patch" ``` -#### 2. Megatron-LM +### 2. Megatron-LM + +Apply the common Megatron patch before the NPU patch. ```bash -export MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 -git clone https://github.com/NVIDIA/Megatron-LM.git "${WORKSPACE}/Megatron-LM" -git -C "${WORKSPACE}/Megatron-LM" checkout "${MEGATRON_COMMIT}" +git clone https://github.com/NVIDIA/Megatron-LM.git "${VIME_INSTALL_ROOT}/Megatron-LM" +git -C "${VIME_INSTALL_ROOT}/Megatron-LM" checkout 1dcf0dafa884ad52ffb243625717a3471643e087 -git -C "${WORKSPACE}/Megatron-LM" apply --whitespace=nowarn "${WORKSPACE}/vime/docker/patch/latest/megatron.patch" -git -C "${WORKSPACE}/Megatron-LM" apply --whitespace=nowarn "${PATCH_DIR}/megatron.patch" +git -C "${VIME_INSTALL_ROOT}/Megatron-LM" apply --check "${VIME_INSTALL_ROOT}/vime/docker/patch/latest/megatron.patch" +git -C "${VIME_INSTALL_ROOT}/Megatron-LM" apply "${VIME_INSTALL_ROOT}/vime/docker/patch/latest/megatron.patch" +git -C "${VIME_INSTALL_ROOT}/Megatron-LM" apply --check "${PATCH_DIR}/megatron.patch" +git -C "${VIME_INSTALL_ROOT}/Megatron-LM" apply "${PATCH_DIR}/megatron.patch" -pip install --no-deps --no-build-isolation -e "${WORKSPACE}/Megatron-LM" +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/Megatron-LM" ``` -#### 3. MegatronAdaptor and TransformerEngineNPU - -The NPU training stack now uses the two source repositories directly. The mainline Megatron patch is applied first; `docker/npu_patch/megatron.patch` contains only the NPU-specific changes rebased onto that mainline patch: +### 3. Megatron-Bridge and mbridge -pip install --no-deps --no-build-isolation -e ${WORKSPACE}/MegatronAdaptor -pip install --no-deps --no-build-isolation -e ${WORKSPACE}/TransformerEngineNPU - -Do not install the CUDA TransformerEngine package in the same environment. - -#### 4. MegatronAdaptor and TransformerEngineNPU +Use the Megatron-Bridge source through `PYTHONPATH`, without installing its CUDA package dependencies. ```bash -export MINDSPEED_COMMIT=fc63de5c48426dd019c3b3f39e65f5bdf56e4086 -git clone https://gitcode.com/Ascend/MindSpeed.git "${WORKSPACE}/MindSpeed" -git -C "${WORKSPACE}/MindSpeed" checkout "${MINDSPEED_COMMIT}" +git clone --branch bridge https://github.com/radixark/Megatron-Bridge.git "${VIME_INSTALL_ROOT}/Megatron-Bridge" +git -C "${VIME_INSTALL_ROOT}/Megatron-Bridge" checkout 3fd3768045422d0aa5c97e90a4e6c659aea9acb9 +git -C "${VIME_INSTALL_ROOT}/Megatron-Bridge" apply --check "${PATCH_DIR}/megatron-bridge.patch" +git -C "${VIME_INSTALL_ROOT}/Megatron-Bridge" apply "${PATCH_DIR}/megatron-bridge.patch" -git -C "${WORKSPACE}/MindSpeed" apply --whitespace=nowarn "${PATCH_DIR}/mindspeed.patch" +git clone https://github.com/ISEEKYAN/mbridge.git "${VIME_INSTALL_ROOT}/mbridge" +git -C "${VIME_INSTALL_ROOT}/mbridge" checkout 89eb10887887bc74853f89a4de258c0702932a1c +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/mbridge" -pip install --no-deps --no-build-isolation -e "${WORKSPACE}/MindSpeed" +pip install --no-build-isolation "nvidia-modelopt==0.46.0" "nvdlfw-inspect==0.2.2" ``` -#### 5. Vime +### 4. TransformerEngineNPU and MegatronAdaptor + +Use TransformerEngineNPU, not the CUDA TransformerEngine package. ```bash -pip install -r "${WORKSPACE}/vime/requirements.txt" -pip install "vllm-router>=0.1.14" -pip install --no-deps --no-build-isolation -e "${WORKSPACE}/vime" +git clone https://gitcode.com/Ascend/TransformerEngineNPU.git "${VIME_INSTALL_ROOT}/TransformerEngineNPU" +git -C "${VIME_INSTALL_ROOT}/TransformerEngineNPU" checkout d743c83d060d5edc48867ecb9e93ec80d81860e4 +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/TransformerEngineNPU" + +git clone https://gitcode.com/Ascend/MegatronAdaptor.git "${VIME_INSTALL_ROOT}/MegatronAdaptor" +git -C "${VIME_INSTALL_ROOT}/MegatronAdaptor" checkout 15582addff3f3d4680e350826fa70d012b475509 +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/MegatronAdaptor" ``` -Build the matching Ascend `torch_memory_saver` wheel. NPU does not actually use -`torch_memory_saver`, but the code still imports and calls it and will break -without it, and there is currently no published Python 3.12 build — so compile -it from source: +### 5. MindSpeed ```bash -git clone --branch 2026.6.0 https://github.com/sgl-project/sgl-kernel-npu.git "${WORKSPACE}/sgl-kernel-npu" -cd "${WORKSPACE}/sgl-kernel-npu" -bash build.sh -a kernels -bash build.sh -a memory-saver -pip install --no-deps output/torch_memory_saver-0.0.8-cp312-cp312-linux_aarch64.whl +git clone https://gitcode.com/Ascend/MindSpeed.git "${VIME_INSTALL_ROOT}/MindSpeed" +git -C "${VIME_INSTALL_ROOT}/MindSpeed" checkout fc63de5c48426dd019c3b3f39e65f5bdf56e4086 +git -C "${VIME_INSTALL_ROOT}/MindSpeed" apply --check "${PATCH_DIR}/mindspeed.patch" +git -C "${VIME_INSTALL_ROOT}/MindSpeed" apply "${PATCH_DIR}/mindspeed.patch" +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/MindSpeed" ``` -#### 5. Install vLLM and vLLM Ascend +### 6. Vime ```bash -export VLLM_COMMIT=9090368b650896bf5fc990c921df7eb4c20355a5 +pip install -r "${VIME_INSTALL_ROOT}/vime/requirements.txt" +pip install --no-deps --no-build-isolation -e "${VIME_INSTALL_ROOT}/vime" +``` -git clone https://github.com/vllm-project/vllm.git "${WORKSPACE}/vllm" -git -C "${WORKSPACE}/vllm" checkout "${VLLM_COMMIT}" -VLLM_TARGET_DEVICE=empty pip install -v -e "${WORKSPACE}/vllm" +### 7. torch_memory_saver -git clone https://github.com/vllm-project/vllm-ascend.git "${WORKSPACE}/vllm-ascend" -git -C "${WORKSPACE}/vllm-ascend" submodule update --init --recursive -pip install -v -e "${WORKSPACE}/vllm-ascend" +Build the NPU wheel from `sgl-kernel-npu`: + +```bash +git clone --depth 1 --branch 2026.6.0 https://github.com/sgl-project/sgl-kernel-npu.git "${VIME_INSTALL_ROOT}/sgl-kernel-npu" +cd "${VIME_INSTALL_ROOT}/sgl-kernel-npu/contrib/torch_memory_saver/python" +python3 setup.py bdist_wheel +python3 -m pip install --no-deps dist/torch_memory_saver-*.whl +cd "${VIME_INSTALL_ROOT}/vime" ``` -> [!NOTE] -> vLLM Ascend has not yet cut a release tag against vLLM 0.22.0. As a temporary -> measure we pin vLLM to the commit below and build vLLM Ascend from source. -> Once vLLM Ascend officially supports 0.22.0, this whole step can be omitted and -> the released packages used instead. +## Environment Setup and Installation Check -## Additional Dependencies +Set the source paths before running Vime: -Ensure the following packages are pinned to these matching versions： +```bash +export PYTHONPATH="${VIME_INSTALL_ROOT}/Megatron-Bridge/src:${VIME_INSTALL_ROOT}/Megatron-LM:${VIME_INSTALL_ROOT}/MegatronAdaptor:${VIME_INSTALL_ROOT}/TransformerEngineNPU:${VIME_INSTALL_ROOT}/vime${PYTHONPATH:+:${PYTHONPATH}}" + +python3 -c 'import megatron, mindspeed, megatron_adaptor, transformer_engine, torch_memory_saver, vime, vllm, vllm_ascend' +``` -```shell -pip install torch-npu==2.10.0.post2 -pip install torchvision==0.25.0 -pip install numpy==1.26.4 -``` \ No newline at end of file +For a complete container build recipe, see [Dockerfile.npu](../Dockerfile.npu). diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch index 0e7798937..adb77c908 100644 --- a/docker/npu_patch/megatron-bridge.patch +++ b/docker/npu_patch/megatron-bridge.patch @@ -1,15 +1,13 @@ diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py -index 63b4cc37..00394727 100644 +index a0421273..7f9203ab 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py -@@ -1199,18 +1199,23 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): - "ColumnParallelLinear", +@@ -1097,15 +1097,20 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): "LinearCrossEntropyModule", "TEColumnParallelLinear", + "MindSpeedTEColumnParallelLinear", "TELayerNormColumnParallelLinear", + "MindSpeedTELayerNormColumnParallelLinear", - "InferenceLayerNormColumnParallelLinear", "TEColumnParallelGroupedLinear", + "MindSpeedTEColumnParallelGroupedLinear", "VocabParallelEmbedding", @@ -20,98 +18,246 @@ index 63b4cc37..00394727 100644 "row": { "RowParallelLinear", "TERowParallelLinear", - "InferenceRowParallelLinear", "TERowParallelGroupedLinear", + "MindSpeedTERowParallelGroupedLinear", }, "replicated": { # Normalization layers -@@ -1282,7 +1287,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): +@@ -1176,7 +1180,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): # Handle fused modules like TELayerNormColumnParallelLinear # These modules have both column-parallel weights (weight, bias) # and replicated layer norm weights (layer_norm_weight, layer_norm_bias) -- if module_type in ("TELayerNormColumnParallelLinear", "InferenceLayerNormColumnParallelLinear"): -+ if module_type in ("TELayerNormColumnParallelLinear", "InferenceLayerNormColumnParallelLinear", "MindSpeedTELayerNormColumnParallelLinear"): +- if module_type == "TELayerNormColumnParallelLinear": ++ if module_type == "TELayerNormColumnParallelLinear" or module_type == "MindSpeedTELayerNormColumnParallelLinear": # Check the actual parameter name to determine the correct parallelism type if self.megatron_param and ( self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias") -@@ -1313,7 +1318,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): +@@ -1207,7 +1211,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): return "replicated" - + # Check parallel_mode for TELinear - if module_type == "TELinear": + if module_type == "TELinear" or module_type == "MindSpeedTELinear": if module.parallel_mode == "column": return "column" elif module.parallel_mode == "row": -diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py -index 7360857d..a6c00358 100644 ---- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py -+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py -@@ -27,6 +27,7 @@ from megatron.core.pipeline_parallel.utils import is_pp_last_stage - from megatron.core.process_groups_config import ProcessGroupCollection - from megatron.core.transformer import MegatronModule +diff --git a/src/megatron/bridge/models/qwen/__init__.py b/src/megatron/bridge/models/qwen/__init__.py +index b3656b6d..382845cc 100644 +--- a/src/megatron/bridge/models/qwen/__init__.py ++++ b/src/megatron/bridge/models/qwen/__init__.py +@@ -15,7 +15,7 @@ + from megatron.bridge.models.qwen.qwen2_bridge import Qwen2Bridge # noqa: F401 + from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge # noqa: F401 + from megatron.bridge.models.qwen.qwen3_moe_bridge import Qwen3MoEBridge # noqa: F401 +-from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge ++# from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge + from megatron.bridge.models.qwen.qwen_provider import ( + Qwen2ModelProvider, + Qwen2ModelProvider1P5B, +@@ -32,8 +32,8 @@ from megatron.bridge.models.qwen.qwen_provider import ( + Qwen3MoEModelProvider, + Qwen3MoEModelProvider30B_A3B, + Qwen3MoEModelProvider235B_A22B, +- Qwen3NextModelProvider, +- Qwen3NextModelProvider80B_A3B, ++ # Qwen3NextModelProvider, ++ # Qwen3NextModelProvider80B_A3B, + Qwen25ModelProvider1P5B, + Qwen25ModelProvider3B, + Qwen25ModelProvider7B, +@@ -67,6 +67,6 @@ __all__ = [ + "Qwen3MoEModelProvider", + "Qwen3MoEModelProvider30B_A3B", + "Qwen3MoEModelProvider235B_A22B", +- "Qwen3NextModelProvider", +- "Qwen3NextModelProvider80B_A3B", ++ # "Qwen3NextModelProvider", ++ # "Qwen3NextModelProvider80B_A3B", + ] +diff --git a/src/megatron/bridge/models/qwen/qwen_provider.py b/src/megatron/bridge/models/qwen/qwen_provider.py +index 775b9765..6200103c 100644 +--- a/src/megatron/bridge/models/qwen/qwen_provider.py ++++ b/src/megatron/bridge/models/qwen/qwen_provider.py +@@ -18,9 +18,9 @@ from typing import TYPE_CHECKING, Callable, Optional + + import torch + import torch.nn.functional as F +-from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( +- get_transformer_block_with_experimental_attention_variant_spec, +-) ++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( ++# get_transformer_block_with_experimental_attention_variant_spec, ++# ) from megatron.core.transformer.spec_utils import ModuleSpec -+from megatron.core.utils import nvtx_range_pop, nvtx_range_push - from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF - - from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention -@@ -497,7 +498,7 @@ class Qwen3VLModel(MegatronModule): - if not _is_mrope_position_ids(position_ids): - position_ids = None - -- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.pre_process") -+ nvtx_range_push(msg="Qwen3VLModel.forward.pre_process") - - cp_rank = self.pg_collection.cp.rank() - cp_size = self.pg_collection.cp.size() -@@ -598,7 +599,7 @@ class Qwen3VLModel(MegatronModule): - vision_embeds, - deepstack_feature_lists, - ) -- torch.cuda.nvtx.range_pop() -+ nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process") - return output_vision_module - else: - vision_embeds = self.vision_embeds -@@ -821,8 +822,8 @@ class Qwen3VLModel(MegatronModule): - if position_ids_were_split and self.language_model is not None: - self.language_model.rotary_pos_emb.is_thd_format = True - -- torch.cuda.nvtx.range_pop() -- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model") -+ nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process") -+ nvtx_range_push(msg="Qwen3VLModel.forward.language_model") - - return_sliced_loss_mask = False - if packed_seq_params is not None: -@@ -863,7 +864,7 @@ class Qwen3VLModel(MegatronModule): - **(extra_block_kwargs or {}), - **kwargs, - ) -- torch.cuda.nvtx.range_pop() -+ nvtx_range_pop(msg="Qwen3VLModel.forward.language_model") - if self.use_dist_train: - if not is_pp_last_stage(self.pg_collection.pp): - return {"language_module": output} + + from megatron.bridge.models.gpt_provider import GPTModelProvider +@@ -430,53 +430,53 @@ class Qwen3MoEModelProvider235B_A22B(Qwen3MoEModelProvider): + # ============================================================================= + + +-@dataclass +-class Qwen3NextModelProvider(Qwen3MoEModelProvider): +- """Base provider for Qwen 3 Next Models.""" +- +- transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( +- get_transformer_block_with_experimental_attention_variant_spec +- ) +- +- layernorm_zero_centered_gamma: bool = True # Zero-centered RMSNorm +- kv_channels: int | None = 256 +- num_query_groups: int = 2 +- seq_length: int = 262144 # 256k tokens +- rotary_base: float = 10000000.0 +- rotary_percent: float = 0.25 # 25% of the hidden size is used for RoPE +- attention_output_gate: bool = True # Gated Attention +- +- # MoE specific parameters +- num_moe_experts: int = 512 +- moe_router_topk: int = 10 # 10 routed experts per token +- moe_shared_expert_gate: bool = True # Qwen3-Next uses a gate for the shared expert +- moe_router_dtype: str = "fp32" +- moe_router_load_balancing_type: str = "global_aux_loss" # Qwen3-Next uses global aux loss for load balancing +- +- # Linear Attention specific parameters +- experimental_attention_variant: str = "gated_delta_net" # Gated Delta Net used in 75% of the model layers +- linear_attention_freq: int | list[int] = 4 # 1 gated standard attention layer per 4 layers +- linear_conv_kernel_dim: int = 4 +- linear_key_head_dim: int = 128 +- linear_value_head_dim: int = 128 +- linear_num_key_heads: int = 16 +- linear_num_value_heads: int = 32 +- +- # Checkpointing +- hetereogenous_dist_checkpoint: bool = True +- +- +-@dataclass +-class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider): +- """ +- Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking +- """ +- +- num_layers: int = 48 +- hidden_size: int = 2048 +- num_attention_heads: int = 16 +- num_query_groups: int = 2 +- ffn_hidden_size: int = 5120 +- moe_ffn_hidden_size: int = 512 +- moe_shared_expert_intermediate_size: int = 512 +- mtp_num_layers: Optional[int] = None ++# @dataclass ++# class Qwen3NextModelProvider(Qwen3MoEModelProvider): ++# """Base provider for Qwen 3 Next Models.""" ++ ++# transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( ++# get_transformer_block_with_experimental_attention_variant_spec ++# ) ++ ++# layernorm_zero_centered_gamma: bool = True # Zero-centered RMSNorm ++# kv_channels: int | None = 256 ++# num_query_groups: int = 2 ++# seq_length: int = 262144 # 256k tokens ++# rotary_base: float = 10000000.0 ++# rotary_percent: float = 0.25 # 25% of the hidden size is used for RoPE ++# attention_output_gate: bool = True # Gated Attention ++ ++# # MoE specific parameters ++# num_moe_experts: int = 512 ++# moe_router_topk: int = 10 # 10 routed experts per token ++# moe_shared_expert_gate: bool = True # Qwen3-Next uses a gate for the shared expert ++# moe_router_dtype: str = "fp32" ++# moe_router_load_balancing_type: str = "global_aux_loss" # Qwen3-Next uses global aux loss for load balancing ++ ++# # Linear Attention specific parameters ++# experimental_attention_variant: str = "gated_delta_net" # Gated Delta Net used in 75% of the model layers ++# linear_attention_freq: int | list[int] = 4 # 1 gated standard attention layer per 4 layers ++# linear_conv_kernel_dim: int = 4 ++# linear_key_head_dim: int = 128 ++# linear_value_head_dim: int = 128 ++# linear_num_key_heads: int = 16 ++# linear_num_value_heads: int = 32 ++ ++# # Checkpointing ++# hetereogenous_dist_checkpoint: bool = True ++ ++ ++# @dataclass ++# class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider): ++# """ ++# Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking ++# """ ++ ++# num_layers: int = 48 ++# hidden_size: int = 2048 ++# num_attention_heads: int = 16 ++# num_query_groups: int = 2 ++# ffn_hidden_size: int = 5120 ++# moe_ffn_hidden_size: int = 512 ++# moe_shared_expert_intermediate_size: int = 512 ++# mtp_num_layers: Optional[int] = None +diff --git a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py +index 6c74827c..75973903 100644 +--- a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py ++++ b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py +@@ -36,9 +36,9 @@ from typing import Any, Callable, List, Optional + + import transformers + from megatron.core.models.gpt import GPTModel as MCoreGPTModel +-from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( +- get_transformer_block_with_experimental_attention_variant_spec, +-) ++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( ++# get_transformer_block_with_experimental_attention_variant_spec, ++# ) + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_block import TransformerBlockSubmodules + from packaging.version import Version as PkgVersion +@@ -105,9 +105,10 @@ class Qwen35VLModelProvider(GPTModelProvider): + # ========================================================================= + # Hybrid Architecture (Qwen3-Next style) + # ========================================================================= +- transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( +- get_transformer_block_with_experimental_attention_variant_spec +- ) ++ # transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( ++ # get_transformer_block_with_experimental_attention_variant_spec ++ # ) ++ transformer_layer_spec: ModuleSpec = None + layernorm_zero_centered_gamma: bool = True + attention_output_gate: bool = True + experimental_attention_variant: str = "gated_delta_net" +@@ -261,9 +262,10 @@ class Qwen35VLMoEModelProvider(GPTModelProvider): + # ========================================================================= + # Hybrid Architecture (Qwen3-Next style) + # ========================================================================= +- transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( +- get_transformer_block_with_experimental_attention_variant_spec +- ) ++ # transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = ( ++ # get_transformer_block_with_experimental_attention_variant_spec ++ # ) ++ transformer_layer_spec: ModuleSpec = None + layernorm_zero_centered_gamma: bool = True + attention_output_gate: bool = True + experimental_attention_variant: str = "gated_delta_net" diff --git a/src/megatron/bridge/models/transformer_config.py b/src/megatron/bridge/models/transformer_config.py -index 6bc5b8d3..dd71b959 100644 +index 618700ed..3133f4ea 100644 --- a/src/megatron/bridge/models/transformer_config.py +++ b/src/megatron/bridge/models/transformer_config.py -@@ -71,6 +71,10 @@ def _resolve_string_fields(config: MCoreTransformerConfig) -> None: - config.pipeline_dtype = str_to_dtype(config.pipeline_dtype) - - +@@ -47,6 +47,10 @@ def _safe_asdict(obj, skip_keys: set[str]) -> dict: + return obj.__class__((_safe_asdict(k, skip_keys), _safe_asdict(v, skip_keys)) for k, v in obj.items()) + return obj + ++from dataclasses import dataclass, field ++ +class MyConfig: + pass -+ -+ + @dataclass class TransformerConfig(MCoreTransformerConfig): - """Megatron Core TransformerConfig with deferred post-init. -@@ -93,6 +97,13 @@ class TransformerConfig(MCoreTransformerConfig): +@@ -70,6 +74,13 @@ class TransformerConfig(MCoreTransformerConfig): """ - + _NO_COPY_KEYS = {"_pg_collection"} + # vllm_eplb_config: MyConfig = field(default_factory=MyConfig) + # vllm_ir_op_priority: MyConfig = field(default_factory=MyConfig) @@ -120,19 +266,110 @@ index 6bc5b8d3..dd71b959 100644 + # vllm_structured_outputs_config: MyConfig = field(default_factory=MyConfig) + # vllm_compilation_config: MyConfig = field(default_factory=MyConfig) + # vllm_attention_config: MyConfig = field(default_factory=MyConfig) - + def __post_init__(self) -> None: """Skip MCore post_init during initial construction. diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py -index d354c34f..d59fd7da 100644 +index 1ca5b18b..3d8ec12a 100644 --- a/src/megatron/bridge/peft/utils.py +++ b/src/megatron/bridge/peft/utils.py -@@ -74,7 +74,7 @@ HAVE_TE = all( +@@ -62,7 +62,7 @@ HAVE_TE = all( ) ) - + -MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") +# MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") - ModelOptLinear, HAVE_MODELOPT_LINEAR = safe_import_from("megatron.core.post_training.modelopt.layers", "Linear") - + TECL = (TEColumnParallelLinear, TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear) + TERL = (TERowParallelLinear, TERowParallelGroupedLinear) +diff --git a/src/megatron/bridge/training/mlm_compat/model.py b/src/megatron/bridge/training/mlm_compat/model.py +index 60cc091c..1dd5c808 100644 +--- a/src/megatron/bridge/training/mlm_compat/model.py ++++ b/src/megatron/bridge/training/mlm_compat/model.py +@@ -21,9 +21,9 @@ from megatron.core import tensor_parallel + from megatron.core.enums import ModelType + from megatron.core.fp8_utils import correct_amax_history_if_needed + from megatron.core.models.gpt import GPTModel +-from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( +- get_transformer_block_with_experimental_attention_variant_spec, +-) ++# from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( ++# get_transformer_block_with_experimental_attention_variant_spec, ++# ) + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_decoder_block_spec, + get_gpt_layer_local_spec, +diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +index 7775c11c..c7b094cf 100644 +--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py ++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +@@ -26,6 +26,7 @@ from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.transformer import MegatronModule + from megatron.core.transformer.spec_utils import ModuleSpec ++from megatron.core.utils import nvtx_range_pop, nvtx_range_push + from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF + + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention +@@ -294,7 +295,7 @@ class Qwen3VLModel(MegatronModule): + # position ids is computed within the model + position_ids = None + +- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.pre_process") ++ nvtx_range_push(msg="Qwen3VLModel.forward.pre_process") + + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() +@@ -387,7 +388,7 @@ class Qwen3VLModel(MegatronModule): + combined_embeddings = split_data_cp_rank(combined_embeddings, cp_size, 0, cp_rank) + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + input_ids_thd, _ = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection + ) +@@ -443,7 +444,7 @@ class Qwen3VLModel(MegatronModule): + # convert lm_input_ids to THD format so it matches position_ids. + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + lm_input_ids, _ = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection + ) +@@ -499,8 +500,8 @@ class Qwen3VLModel(MegatronModule): + attention_mask = None + self.language_model.rotary_pos_emb.is_thd_format = True + +- torch.cuda.nvtx.range_pop() +- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model") ++ nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process") ++ nvtx_range_push(msg="Qwen3VLModel.forward.language_model") + + output = self.language_model( + input_ids=lm_input_ids, +@@ -516,6 +517,6 @@ class Qwen3VLModel(MegatronModule): + **(extra_block_kwargs or {}), + **kwargs, + ) +- torch.cuda.nvtx.range_pop() ++ nvtx_range_pop(msg="Qwen3VLModel.forward.language_model") + + return output +diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +index b714d0f7..42fcab74 100644 +--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py ++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +@@ -668,6 +668,11 @@ def preprocess_packed_seqs( + """ + batch_size = input_ids.shape[0] + ++ # Ensure boolean dtype for correct advanced indexing (bool → mask select, ++ # int → fancy index which silently corrupts data when values are 0/1). ++ if attention_mask.dtype != torch.bool: ++ attention_mask = attention_mask.bool() ++ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + if pg_collection is not None: + tp_size = pg_collection.tp.size() diff --git a/docker/npu_patch/megatron.patch b/docker/npu_patch/megatron.patch index 7ac850d96..e676ac3f6 100644 --- a/docker/npu_patch/megatron.patch +++ b/docker/npu_patch/megatron.patch @@ -455,4906 +455,29 @@ index dd8590a79..3736b0e31 100644 assert len(steps) == 1 step = torch.tensor(steps[0], dtype=torch.float) -diff --git a/megatron/core/ssm/chunk_gated_delta_rule.py b/megatron/core/ssm/chunk_gated_delta_rule.py -new file mode 100644 -index 000000000..dc388fd7f ---- /dev/null -+++ b/megatron/core/ssm/chunk_gated_delta_rule.py -@@ -0,0 +1,490 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+import warnings -+from typing import Optional -+ -+import torch -+ -+from .triton.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -+from .triton.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o -+from .triton.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -+from .triton.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd -+from .triton.solve_tril import solve_tril -+from .triton.cumsum import chunk_local_cumsum -+from .triton.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard -+from .triton.l2norm import l2norm -+ -+import torch_npu -+import fla_npu -+ -+ -+def prepare_chunk_indices( -+ cu_seqlens: list[int], -+ chunk_size: int -+) -> list[int]: -+ """ -+ Generate chunk indices based on cu_seqlens (list[int]). -+ -+ Note: The original PyTorch version returns a Tensor of shape [N, 2]. -+ To maintain pure Python compatibility, this version returns list[tuple[start_seq_idx, chunk_idx_in_seq]]. -+ If the operator requires a flattened list[int] (e.g., [s0, c0, s1, c1, ...]), please flatten it before calling. -+ -+ Logic replication from original code: -+ 1. Calculate length for each sequence: lens[i] = cu_seqlens[i+1] - cu_seqlens[i] -+ 2. Calculate number of chunks needed for each sequence: ceil(lens[i] / chunk_size) -+ 3. Generate corresponding (sequence_id, chunk_id) pairs -+ """ -+ indices = [] -+ -+ # Iterate over each sequence segment -+ for i in range(len(cu_seqlens) - 1): -+ start = cu_seqlens[i] -+ end = cu_seqlens[i+1] -+ length = end - start -+ -+ if length <= 0: -+ continue -+ -+ # Calculate how many chunks are needed for this sequence -+ # Equivalent to cdiv(length, chunk_size) -+ num_chunks = (length + chunk_size - 1) // chunk_size -+ -+ for chunk_id in range(num_chunks): -+ # Original logic: indices.eq(0).cumsum(0) - 1 corresponds to sequence index i -+ # Original logic: indices corresponds to chunk_id -+ indices.append((i)) -+ indices.append((chunk_id)) -+ -+ return indices -+ -+ -+def chunk_gated_delta_rule_fwd( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ g: torch.Tensor, -+ beta: torch.Tensor, -+ scale: float, -+ initial_state: torch.Tensor, -+ output_final_state: bool, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64, -+): -+ g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, head_first=False) -+ # obtain WY representation. u is actually the new v. -+ A = chunk_scaled_dot_kkt_fwd( -+ k=k, -+ g=g, -+ beta=beta, -+ cu_seqlens=cu_seqlens, -+ chunk_size=chunk_size, -+ output_dtype=torch.float32 -+ ) -+ A = solve_tril( -+ A=A, -+ cu_seqlens=cu_seqlens, -+ output_dtype=k.dtype -+ ) -+ -+ q = q.transpose(1, 2).contiguous() -+ k = k.transpose(1, 2).contiguous() -+ v = v.transpose(1, 2).contiguous() -+ g = g.transpose(1, 2).contiguous() -+ A = A.transpose(1, 2).contiguous() -+ beta = beta.transpose(1, 2).contiguous().float() -+ -+ if cu_seqlens is not None: -+ cu_seqlens1 = cu_seqlens.tolist() -+ chunk_indices = prepare_chunk_indices(cu_seqlens1, chunk_size) -+ else: -+ cu_seqlens1 = cu_seqlens -+ chunk_indices = None -+ -+ w, u = torch.ops.npu.npu_recompute_w_u_fwd( -+ k, -+ v, -+ beta, -+ A, -+ chunk_size, -+ g = g, -+ gk = None, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices -+ ) -+ -+ h, v_new, final_state = torch.ops.npu.npu_chunk_gated_delta_rule_fwd_h( -+ k, -+ w, -+ u, -+ g=g, -+ initial_state=initial_state, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ output_final_state=output_final_state, -+ chunk_size=chunk_size -+ ) -+ -+ o = torch.ops.npu.npu_chunk_fwd_o( -+ q, -+ k, -+ v_new, -+ h, -+ scale, -+ g=g, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ chunk_size=chunk_size -+ ) -+ -+ g = g.transpose(1, 2).contiguous() -+ o = o.transpose(1, 2).contiguous() -+ -+ return g, o, A, final_state -+ -+ -+def chunk_gated_delta_rule_bwd( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ g: torch.Tensor, -+ beta: torch.Tensor, -+ A: torch.Tensor, -+ scale: float, -+ initial_state: torch.Tensor, -+ do: torch.Tensor, -+ dht: torch.Tensor, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64, -+): -+ q = q.transpose(1, 2).contiguous() -+ k = k.transpose(1, 2).contiguous() -+ v = v.transpose(1, 2).contiguous() -+ g = g.transpose(1, 2).contiguous() -+ beta = beta.transpose(1, 2).contiguous().float() -+ do = do.transpose(1, 2).contiguous() -+ # Note: A is not transposed here because it was already transposed in the forward pass -+ -+ if cu_seqlens is not None: -+ cu_seqlens1 = cu_seqlens.tolist() -+ chunk_indices = prepare_chunk_indices(cu_seqlens1, chunk_size) -+ else: -+ cu_seqlens1 = cu_seqlens -+ chunk_indices = None -+ -+ w, u = torch.ops.npu.npu_recompute_w_u_fwd( -+ k, -+ v, -+ beta, -+ A, -+ chunk_size, -+ g = g, -+ gk = None, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices -+ ) -+ -+ h, v_new, final_state = torch.ops.npu.npu_chunk_gated_delta_rule_fwd_h( -+ k, -+ w, -+ u, -+ g, -+ initial_state=initial_state, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ output_final_state=False, -+ chunk_size=chunk_size -+ ) -+ -+ dv = torch.ops.npu.npu_chunk_bwd_dv_local( -+ q, -+ k, -+ do, -+ g, -+ g_gamma=None, -+ A=A, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ scale=scale, -+ chunk_size=chunk_size -+ ) -+ -+ dh, dh0, dv = torch.ops.npu.npu_chunk_gated_delta_rule_bwd_dhu( -+ q, -+ k, -+ w, -+ do, -+ dv, -+ g=g, -+ gK=None, -+ h0=None, -+ dht=dht, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ scale=scale, -+ chunk_size=chunk_size -+ ) -+ -+ dq, dk, dw, dg = torch.ops.npu.npu_chunk_bwd_dqkwg( -+ q, -+ k, -+ v_new, -+ g, -+ h, -+ do, -+ dh, -+ dv, -+ chunk_size, -+ chunk_indices=chunk_indices, -+ scale=scale, -+ cu_seqlens=cu_seqlens1 -+ ) -+ -+ dq = dq.transpose(1, 2).contiguous() -+ dk = dk.transpose(1, 2).contiguous() -+ dg = dg.transpose(1, 2).contiguous() -+ -+ dA = torch.ops.npu.npu_prepare_wy_repr_bwd_da( -+ k, -+ v, -+ beta, -+ A, -+ dw, -+ dv, -+ g, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ chunk_size=chunk_size -+ ) -+ -+ dk2, dv, db, dg2 = torch.ops.npu.npu_prepare_wy_repr_bwd_full( -+ k, -+ v, -+ beta, -+ A, -+ dA, -+ dw, -+ dv, -+ g, -+ chunk_size, -+ cu_seqlens=cu_seqlens1, -+ chunk_indices=chunk_indices, -+ ) -+ dk2 = dk2.transpose(1, 2).contiguous() -+ dv = dv.transpose(1, 2).contiguous() -+ db = db.transpose(1, 2).contiguous() -+ dg2 = dg2.transpose(1, 2).contiguous() -+ -+ dk.add_(dk2) -+ dg.add_(dg2) -+ if dg.dtype != torch.float32: -+ raise ValueError( -+ f"dg current type is {dg.dtype} , should be float32" -+ ) -+ -+ dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, head_first=False) -+ -+ return dq, dk, dv, db, dg, dh0 -+ -+ -+class ChunkGatedDeltaRuleFunction(torch.autograd.Function): -+ -+ @staticmethod -+ @input_guard -+ @autocast_custom_fwd -+ def forward( -+ ctx, -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ g: torch.Tensor, -+ beta: torch.Tensor, -+ scale: float, -+ initial_state: torch.Tensor, -+ output_final_state: bool, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ use_qk_l2norm_in_kernel: bool = False, -+ chunk_size: int = 64, -+ ): -+ q_rstd, k_rstd = None, None -+ g, o, A, final_state = chunk_gated_delta_rule_fwd( -+ q=q, -+ k=k, -+ v=v, -+ g=g, -+ beta=beta, -+ scale=scale, -+ initial_state=initial_state, -+ output_final_state=output_final_state, -+ cu_seqlens=cu_seqlens, -+ chunk_size=chunk_size -+ ) -+ ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens) -+ ctx.scale = scale -+ ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel -+ ctx.chunk_size = chunk_size -+ return o.to(q.dtype), final_state -+ -+ @staticmethod -+ @input_guard -+ @autocast_custom_bwd -+ def backward( -+ ctx, -+ do: torch.Tensor, -+ dht: torch.Tensor -+ ): -+ q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens = ctx.saved_tensors -+ dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd( -+ q=q, -+ k=k, -+ v=v, -+ g=g, -+ beta=beta, -+ A=A, -+ scale=ctx.scale, -+ initial_state=initial_state, -+ do=do, -+ dht=dht, -+ cu_seqlens=cu_seqlens, -+ chunk_size=ctx.chunk_size, -+ ) -+ return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None -+ -+ -+@torch.compiler.disable -+def chunk_gated_delta_rule( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ g: torch.Tensor, -+ beta: torch.Tensor, -+ scale: float = None, -+ initial_state: torch.Tensor = None, -+ output_final_state: bool = False, -+ use_qk_l2norm_in_kernel: bool = False, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64, -+ head_first: bool = False, -+): -+ r""" -+ Args: -+ q (torch.Tensor): -+ queries of shape `[B, T, H, K]`. -+ k (torch.Tensor): -+ keys of shape `[B, T, H, K]`. -+ v (torch.Tensor): -+ values of shape `[B, T, H, V]`. -+ g (torch.Tensor): -+ (forget) gating tensor (in log space!) of shape `[B, T, H]`. -+ beta (torch.Tensor): -+ betas of shape `[B, T, H]`. -+ scale (Optional[float]): -+ Scale factor for the RetNet attention scores. -+ If not provided, it will default to `1 / sqrt(K)`. Default: `None`. -+ initial_state (Optional[torch.Tensor]): -+ Initial state of shape `[N, H, K, V]` for `N` input sequences. -+ For equal-length input sequences, `N` equals the batch size `B`. -+ Default: `None`. -+ output_final_state (Optional[bool]): -+ Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. -+ use_qk_l2norm_in_kernel (bool): -+ Whether to apply L2norm to the q/k tensor internally. Default: `False`. -+ cu_seqlens (torch.LongTensor): -+ Cumulative sequence lengths of shape `[N+1]` used for variable-length training, -+ consistent with the FlashAttention API. -+ head_first (Optional[bool]): -+ Whether the inputs are in the head-first format. Default: `False`. -+ This argument has been deprecated. -+ -+ Returns: -+ o (torch.Tensor): -+ Outputs of shape `[B, T, H, V]`. -+ final_state (torch.Tensor): -+ Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. -+ -+ Examples:: -+ >>> import torch -+ >>> import torch.nn.functional as F -+ >>> from einops import rearrange -+ >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule -+ # inputs with equal lengths -+ >>> B, T, H, K, V = 4, 2048, 4, 512, 512 -+ >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') -+ >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) -+ >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') -+ >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() -+ >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) -+ >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') -+ >>> o, ht = chunk_gated_delta_rule( -+ ... q, k, v, g, beta, -+ ... initial_state=h0, -+ ... output_final_state=True -+ ... ) -+ # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required -+ >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) -+ # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected -+ >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) -+ >>> o, ht = chunk_gated_delta_rule( -+ ... q, k, v, g, beta, -+ ... initial_state=h0, -+ ... output_final_state=True, -+ ... cu_seqlens=cu_seqlens -+ ... ) -+ """ -+ if q.dtype != k.dtype or k.dtype != v.dtype: -+ raise ValueError( -+ f"q current type is {q.dtype} , k current type is {k.dtype} ,v current type is {v.dtype} , they should are equal" -+ ) -+ if q.dtype == torch.float32: -+ raise ValueError( -+ "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." -+ ) -+ if len(beta.shape) != 3: -+ raise ValueError( -+ f"beta current shape len is {len(beta.shape)}, beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." -+ ) -+ -+ if head_first: -+ warnings.warn( -+ "head_first is deprecated and will be removed in a future version. " -+ "Please use head_first=False for now instead." -+ ) -+ if not head_first and q.shape[1] < q.shape[2]: -+ warnings.warn( -+ f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). " -+ "This may indicate the inputs were passed in head-first format [B, H, T, ...] " -+ "when head_first=False was specified. " -+ "Please verify your input tensor format matches the expected shape [B, T, H, ...]." -+ ) -+ if cu_seqlens is not None: -+ if q.shape[0] != 1: -+ raise ValueError( -+ f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." -+ f"Please flatten variable-length inputs before processing." -+ ) -+ if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: -+ raise ValueError( -+ f"The number of initial states is expected to be equal to the number of input sequences, " -+ f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." -+ ) -+ if scale is None: -+ scale = k.shape[-1] ** -0.5 -+ -+ if use_qk_l2norm_in_kernel: -+ q = l2norm(q, dim=-1, eps=1e-6) -+ k = l2norm(k, dim=-1, eps=1e-6) -+ -+ o, final_state = ChunkGatedDeltaRuleFunction.apply( -+ q, -+ k, -+ v, -+ g, -+ beta, -+ scale, -+ initial_state, -+ output_final_state, -+ cu_seqlens, -+ use_qk_l2norm_in_kernel, -+ chunk_size -+ ) -+ return o, final_state -diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py -index 601a72a43..8998e70f9 100644 ---- a/megatron/core/ssm/gated_delta_net.py -+++ b/megatron/core/ssm/gated_delta_net.py -@@ -40,9 +40,9 @@ from megatron.core.transformer.utils import ( - from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push - - try: -- from fla.modules.convolution import causal_conv1d -- from fla.modules.l2norm import l2norm -- from fla.ops.gated_delta_rule import chunk_gated_delta_rule -+ from megatron.core.ssm.triton.causal_conv1d import causal_conv1d -+ from megatron.core.ssm.triton.l2norm import l2norm -+ from megatron.core.ssm.chunk_gated_delta_rule import chunk_gated_delta_rule - - HAVE_FLA = True - except ImportError: -@@ -377,14 +377,13 @@ class GatedDeltaNet(MegatronModule): - qkv = self.act_fn(conv_out[..., :seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d - else: -- assert self.activation in ["silu", "swish"] -+ conv1d_weight = conv1d_weight.squeeze(1) - qkv, _ = causal_conv1d( -- x=qkv, # FLA conv1d accepts [b, s, d] format input -- weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w -+ x=qkv, -+ weight=conv1d_weight.transpose(-1, -2).contiguous(), - bias=conv1d_bias, -- activation=self.activation, -- initial_state=None, -- output_final_state=False, -+ activation='silu', -+ cu_seqlens=None - ) - nvtx_range_pop(suffix="conv1d") - -@@ -465,7 +464,6 @@ class GatedDeltaNet(MegatronModule): - - return out, out_bias - -- @jit_fuser - def _apply_gated_norm(self, x, gate): - # Output Norm - x_dtype = x.dtype -diff --git a/megatron/core/ssm/triton/causal_conv1d.py b/megatron/core/ssm/triton/causal_conv1d.py -new file mode 100644 -index 000000000..c430c82aa ---- /dev/null -+++ b/megatron/core/ssm/triton/causal_conv1d.py -@@ -0,0 +1,120 @@ -+# Copyright (c) 2026, Huawei Technologies Co., Ltd. -+# Copyright 2025 XPU-Forces Team -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+ -+from typing import Optional -+ -+import torch -+ -+from .convolution import ( -+ causal_conv1d_fwd_impl, -+ causal_conv1d_bwd_impl, -+) -+ -+from .utils import is_arch35 -+ -+__all__ = ["CausalConv1dFunction", "causal_conv1d"] -+ -+ -+# Placeholder used in ctx.save_for_backward since it does not accept None -+_PLACEHOLDER = torch.empty(0) -+ -+ -+class CausalConv1dFunction(torch.autograd.Function): -+ @staticmethod -+ def forward( -+ ctx, -+ x: torch.Tensor, -+ weight: torch.Tensor, -+ bias: Optional[torch.Tensor] = None, -+ residual: Optional[torch.Tensor] = None, -+ initial_state: Optional[torch.Tensor] = None, -+ activation: str = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ output_final_state: bool = False, -+ ): -+ if is_arch35(): -+ raise NotImplementedError("causal_conv1d is not supported on arch35") -+ -+ y, final_state = causal_conv1d_fwd_impl( -+ x=x, -+ weight=weight, -+ bias=bias, -+ residual=residual, -+ initial_state=initial_state, -+ activation=activation, -+ cu_seqlens=cu_seqlens, -+ output_final_state=output_final_state, -+ ) -+ -+ # save_for_backward does not accept None — use _PLACEHOLDER instead -+ ctx.save_for_backward( -+ x, -+ weight, -+ bias if bias is not None else _PLACEHOLDER, -+ residual if residual is not None else _PLACEHOLDER, -+ initial_state if initial_state is not None else _PLACEHOLDER, -+ cu_seqlens if cu_seqlens is not None else _PLACEHOLDER, -+ ) -+ ctx.has_bias = bias is not None -+ ctx.has_residual = residual is not None -+ ctx.has_initial_state = initial_state is not None -+ ctx.has_cu_seqlens = cu_seqlens is not None -+ ctx.activation = activation -+ -+ return y, final_state -+ -+ @staticmethod -+ def backward(ctx, dy: torch.Tensor, d_final_state: Optional[torch.Tensor] = None): -+ if is_arch35(): -+ raise NotImplementedError("causal_conv1d is not supported on arch35") -+ -+ x, weight, bias, residual, initial_state, cu_seqlens = ctx.saved_tensors -+ -+ # Restore None placeholders -+ bias = bias if ctx.has_bias else None -+ residual = residual if ctx.has_residual else None -+ initial_state = initial_state if ctx.has_initial_state else None -+ cu_seqlens = cu_seqlens if ctx.has_cu_seqlens else None -+ -+ # bwd_impl has @input_guard(make_contiguous=True); no manual handling needed -+ dx, dw, db, dr, dh0 = causal_conv1d_bwd_impl( -+ x=x, -+ dy=dy, -+ dht=d_final_state, -+ weight=weight, -+ bias=bias, -+ residual=residual, -+ initial_state=initial_state, -+ activation=ctx.activation, -+ cu_seqlens=cu_seqlens, -+ ) -+ -+ # Return order must match forward args: -+ # x, weight, bias, residual, initial_state, activation, cu_seqlens, output_final_state -+ return dx, dw, db, dr, dh0, None, None, None -+ -+ -+def causal_conv1d( -+ x: torch.Tensor, -+ weight: torch.Tensor, -+ bias: Optional[torch.Tensor] = None, -+ residual: Optional[torch.Tensor] = None, -+ initial_state: Optional[torch.Tensor] = None, -+ activation: str = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ output_final_state: bool = False, -+): -+ return CausalConv1dFunction.apply( -+ x, -+ weight, -+ bias, -+ residual, -+ initial_state, -+ activation, -+ cu_seqlens, -+ output_final_state, -+ ) -+ -diff --git a/megatron/core/ssm/triton/chunk_delta_h.py b/megatron/core/ssm/triton/chunk_delta_h.py -new file mode 100644 -index 000000000..4dfc37f20 ---- /dev/null -+++ b/megatron/core/ssm/triton/chunk_delta_h.py -@@ -0,0 +1,572 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+from typing import Optional, Tuple -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices, prepare_chunk_offsets, get_autotune_config, get_npu_properties -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_GK': lambda args: args['gk'] is not None, -+ 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, -+ 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, -+ 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -+}) -+@triton.autotune( -+ configs=get_autotune_config(multibuffer_list=(False,)), -+ key=['H', 'K', 'V', 'BT'], -+) -+@triton.jit(do_not_specialize=['T']) -+def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( -+ k, -+ v, -+ w, -+ v_new, -+ g, -+ gk, -+ h, -+ h0, -+ ht, -+ cu_seqlens, -+ chunk_offsets, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ BT: tl.constexpr, -+ BV: tl.constexpr, -+ NT: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_GK: tl.constexpr, -+ USE_INITIAL_STATE: tl.constexpr, -+ STORE_FINAL_STATE: tl.constexpr, -+ SAVE_NEW_VALUE: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ T_all = T -+ NT_all = NT -+ i_v, i_nh = tl.program_id(0), tl.program_id(1) -+ i_n, i_h = i_nh // H, i_nh % H -+ if IS_VARLEN: -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T = eos - bos -+ NT = tl.cdiv(T, BT) -+ boh = tl.load(chunk_offsets + i_n).to(tl.int32) -+ else: -+ bos, eos = i_n * T, i_n * T + T -+ NT = tl.cdiv(T, BT) -+ boh = i_n * NT -+ -+ # Initialize hidden states -+ b_h1 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 64: -+ b_h2 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 128: -+ b_h3 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 192: -+ b_h4 = tl.zeros([64, BV], dtype=tl.float32) -+ -+ if IS_VARLEN: -+ v = v + (i_h * T_all + bos) * V -+ k = k + (i_h * T_all + bos) * K -+ w = w + (i_h * T_all + bos) * K -+ g = g + i_h * T_all + bos -+ h = h + (i_h * NT_all + boh) * K * V -+ if SAVE_NEW_VALUE: -+ v_new_base = v_new + (i_h * T_all + bos) * V -+ else: -+ v = v + (i_n * H + i_h) * T * V -+ k = k + (i_n * H + i_h) * T * K -+ w = w + (i_n * H + i_h) * T * K -+ g = g + (i_n * H + i_h) * T -+ h = h + (i_n * H + i_h) * NT * K * V -+ if SAVE_NEW_VALUE: -+ v_new_base = v_new + (i_n * H + i_h) * T * V -+ -+ if USE_INITIAL_STATE: -+ h0_ptr = h0 + i_nh * K * V -+ if STORE_FINAL_STATE: -+ ht_ptr = ht + i_nh * K * V -+ -+ # Load initial state -+ if USE_INITIAL_STATE: -+ p_h0_1 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) -+ if K > 64: -+ p_h0_2 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) -+ if K > 128: -+ p_h0_3 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) -+ if K > 192: -+ p_h0_4 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) -+ -+ # Main recurrence over chunks -+ for i_t in range(NT): -+ # Store current hidden state h_t -+ p_h1 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 64: -+ p_h2 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 128: -+ p_h3 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 192: -+ p_h4 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) -+ -+ # Compute v_residual = v - w @ h -+ p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) -+ if K > 64: -+ p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) -+ if K > 128: -+ p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) -+ if K > 192: -+ p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) -+ -+ p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v -+ -+ if SAVE_NEW_VALUE: -+ p_v_new = tl.make_block_ptr(v_new_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) -+ -+ last_idx = min((i_t + 1) * BT, T) - 1 -+ -+ # Apply output gate g -+ if USE_G: -+ m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T -+ b_g_last = tl.load(g + last_idx) -+ p_g = tl.make_block_ptr(g, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_v *= (m_t * tl.exp(b_g_last - b_g))[:, None] -+ b_g_last_exp = tl.exp(b_g_last) -+ b_h1 *= b_g_last_exp -+ if K > 64: -+ b_h2 *= b_g_last_exp -+ if K > 128: -+ b_h3 *= b_g_last_exp -+ if K > 192: -+ b_h4 *= b_g_last_exp -+ -+ # Apply key gate gk -+ if USE_GK: -+ o_k1 = tl.arange(0, 64).to(tl.float32) -+ gk_base_ptr = gk + (i_n * H + i_h) * T * K -+ b_gk_last1 = tl.load(gk_base_ptr + last_idx * K + o_k1, mask=(o_k1 < K), other=0.) -+ b_h1 *= tl.exp(b_gk_last1)[:, None] -+ if K > 64: -+ o_k2 = 64 + o_k1 -+ b_gk_last2 = tl.load(gk_base_ptr + last_idx * K + o_k2, mask=(o_k2 < K), other=0.) -+ b_h2 *= tl.exp(b_gk_last2)[:, None] -+ if K > 128: -+ o_k3 = 128 + o_k1 -+ b_gk_last3 = tl.load(gk_base_ptr + last_idx * K + o_k3, mask=(o_k3 < K), other=0.) -+ b_h3 *= tl.exp(b_gk_last3)[:, None] -+ if K > 192: -+ o_k4 = 192 + o_k1 -+ b_gk_last4 = tl.load(gk_base_ptr + last_idx * K + o_k4, mask=(o_k4 < K), other=0.) -+ b_h4 *= tl.exp(b_gk_last4)[:, None] -+ -+ b_v = b_v.to(k.dtype.element_ty) -+ -+ # Update hidden state: h += k @ v -+ p_k = tl.make_block_ptr(k, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1)) -+ b_k = (b_k * tl.exp(b_gk_last1[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) -+ b_h1 += tl.dot(b_k, b_v) -+ -+ if K > 64: -+ p_k = tl.make_block_ptr(k, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1)) -+ b_k = (b_k * tl.exp(b_gk_last2[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) -+ b_h2 += tl.dot(b_k, b_v) -+ -+ if K > 128: -+ p_k = tl.make_block_ptr(k, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1)) -+ b_k = (b_k * tl.exp(b_gk_last3[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) -+ b_h3 += tl.dot(b_k, b_v) -+ -+ if K > 192: -+ p_k = tl.make_block_ptr(k, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1)) -+ b_k = (b_k * tl.exp(b_gk_last4[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype) -+ b_h4 += tl.dot(b_k, b_v) -+ -+ # Store final state -+ if STORE_FINAL_STATE: -+ p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 64: -+ p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 128: -+ p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 192: -+ p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+def chunk_gated_delta_rule_fwd_h( -+ k: torch.Tensor, -+ w: torch.Tensor, -+ u: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ gk: Optional[torch.Tensor] = None, -+ initial_state: Optional[torch.Tensor] = None, -+ output_final_state: bool = False, -+ chunk_size: int = 64, # default:64 -+ save_new_value: bool = True, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+) -> Tuple[torch.Tensor, torch.Tensor]: -+ B, T, H, K, V = *k.shape, u.shape[-1] -+ BT = chunk_size -+ -+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None -+ # N: the actual number of sequences in the batch with either equal or variable lengths -+ if cu_seqlens is None: -+ N, NT, chunk_offsets = B, triton.cdiv(T, BT), None -+ else: -+ N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) -+ assert K <= 256, "current kernel does not support head dimension larger than 256." -+ -+ h = k.new_empty(B, NT, H, K, V).permute(0, 2, 1, 3, 4).contiguous() -+ final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None -+ -+ BV = 128 -+ -+ v_new = torch.empty_like(u).permute(0, 2, 1, 3).contiguous() if save_new_value else None -+ k = k.permute(0, 2, 1, 3).contiguous() -+ w = w.permute(0, 2, 1, 3).contiguous() -+ u = u.permute(0, 2, 1, 3).contiguous() -+ g = g.permute(0, 2, 1).contiguous() -+ chunk_gated_delta_rule_fwd_kernel_h_blockdim64[(triton.cdiv(V, BV), N * H)]( -+ k=k, -+ v=u, -+ w=w, -+ v_new=v_new, -+ g=g, -+ gk=gk, -+ h=h, -+ h0=initial_state, -+ ht=final_state, -+ cu_seqlens=cu_seqlens, -+ chunk_offsets=chunk_offsets, -+ T=T, -+ H=H, -+ K=K, -+ V=V, -+ BT=BT, -+ BV=BV, -+ NT=NT, -+ ) -+ h = h.permute(0, 2, 1, 3, 4).contiguous() -+ v_new = v_new.permute(0, 2, 1, 3).contiguous() -+ return h, v_new, final_state -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_GK': lambda args: args['gk'] is not None, -+ 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, -+ 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -+}) -+@triton.autotune( -+ configs=get_autotune_config(multibuffer_list=(True, False)), -+ key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'IS_VARLEN'], -+) -+@triton.jit(do_not_specialize=['T']) -+def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( -+ q, -+ k, -+ w, -+ g, -+ gk, -+ dht, -+ dh0, -+ do, -+ dh, -+ dv, -+ dv2, -+ cu_seqlens, -+ chunk_offsets, -+ scale, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ BT: tl.constexpr, -+ BV: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_GK: tl.constexpr, -+ USE_INITIAL_STATE: tl.constexpr, -+ USE_FINAL_STATE_GRADIENT: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ T_all = T -+ i_v, i_nh = tl.program_id(0), tl.program_id(1) -+ i_n, i_h = i_nh // H, i_nh % H -+ if IS_VARLEN: -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T = eos - bos -+ NT = tl.cdiv(T, BT) -+ boh = tl.load(chunk_offsets + i_n).to(tl.int32) -+ else: -+ bos, eos = i_n * T, i_n * T + T -+ NT = tl.cdiv(T, BT) -+ boh = i_n * NT -+ -+ b_dh1 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 64: -+ b_dh2 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 128: -+ b_dh3 = tl.zeros([64, BV], dtype=tl.float32) -+ if K > 192: -+ b_dh4 = tl.zeros([64, BV], dtype=tl.float32) -+ -+ q += (bos * H + i_h) * K -+ k += (bos * H + i_h) * K -+ w += (bos * H + i_h) * K -+ do += (bos * H + i_h) * V -+ dv += (bos * H + i_h) * V -+ dv2 += (bos * H + i_h) * V -+ dh += (boh * H + i_h) * K * V -+ if USE_GK: -+ gk += (bos * H + i_h) * K -+ -+ if USE_INITIAL_STATE: -+ dh0 += i_nh * K * V -+ if USE_FINAL_STATE_GRADIENT: -+ dht += i_nh * K * V -+ -+ stride_v = H * V -+ stride_h = H * K * V -+ stride_k = H * K -+ -+ if USE_FINAL_STATE_GRADIENT: -+ p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) -+ if K > 64: -+ p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) -+ if K > 128: -+ p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) -+ if K > 192: -+ p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) -+ -+ for i_t in range(NT - 1, -1, -1): -+ p_dh1 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 64: -+ p_dh2 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 128: -+ p_dh3 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 192: -+ p_dh4 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) -+ -+ last_idx = min((i_t + 1) * BT, T) - 1 -+ if USE_G: -+ if IS_VARLEN: -+ bos_g = i_h * T_all + bos -+ else: -+ bos_g = (i_n * H + i_h) * T_all -+ bg_last = tl.load(g + bos_g + last_idx) -+ bg_last_exp = tl.exp(bg_last) -+ p_g = tl.make_block_ptr(base=g + bos_g, shape=(T,), strides=(1,), offsets=(i_t * BT,), block_shape=(BT,), order=(0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_g_exp = tl.exp(b_g) -+ -+ p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ -+ b_do = tl.load(p_do, boundary_check=(0, 1)) -+ -+ # Update dv -+ p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ o_k1 = tl.arange(0, 64) -+ b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.) -+ b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) -+ -+ if K > 64: -+ p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ o_k2 = 64 + o_k1 -+ b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.) -+ b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) -+ -+ if K > 128: -+ p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ o_k3 = 128 + o_k1 -+ b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.) -+ b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) -+ -+ if K > 192: -+ p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ if USE_GK: -+ o_k4 = 192 + o_k1 -+ b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.) -+ b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) -+ -+ if USE_G: -+ m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T -+ b_dv *= (m_t * tl.exp(bg_last - b_g))[:, None] -+ b_dv += tl.load(p_dv, boundary_check=(0, 1)) -+ -+ tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) -+ # Update dh -+ p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) -+ p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ if USE_G: -+ b_dh1 *= bg_last_exp -+ b_q = b_q * b_g_exp[None, :] -+ if USE_GK: -+ b_dh1 *= tl.exp(b_gk_last1[:, None]) -+ b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) -+ if K > 64: -+ p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) -+ p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ if USE_G: -+ b_dh2 *= bg_last_exp -+ b_q = b_q * b_g_exp[None, :] -+ if USE_GK: -+ b_dh2 *= tl.exp(b_gk_last2[:, None]) -+ b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) -+ if K > 128: -+ p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) -+ p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ if USE_G: -+ b_dh3 *= bg_last_exp -+ b_q = b_q * b_g_exp[None, :] -+ if USE_GK: -+ b_dh3 *= tl.exp(b_gk_last3[:, None]) -+ b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) -+ if K > 192: -+ p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) -+ p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ if USE_G: -+ b_dh4 *= bg_last_exp -+ b_q = b_q * b_g_exp[None, :] -+ if USE_GK: -+ b_dh4 *= tl.exp(b_gk_last4[:, None]) -+ b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) -+ -+ if USE_INITIAL_STATE: -+ p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 64: -+ p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 128: -+ p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) -+ if K > 192: -+ p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) -+ tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+def chunk_gated_delta_rule_bwd_dhu( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ w: torch.Tensor, -+ do: torch.Tensor, -+ dv: torch.Tensor, -+ g: torch.Tensor | None = None, -+ gk: torch.Tensor | None = None, -+ h0: torch.Tensor | None = None, -+ dht: torch.Tensor | None = None, -+ scale: float | None = None, -+ cu_seqlens: torch.LongTensor | None = None, -+ chunk_size: int = 64, # SY: remove this argument and force chunk size 64? -+ chunk_indices: torch.LongTensor | None = None, -+ use_exp2: bool = False, -+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: -+ B, T, H, K, V = *q.shape, do.shape[-1] -+ # N: the actual number of sequences in the batch with either equal or variable lengths -+ BT = 64 -+ assert K <= 256, "current kernel does not support head dimension being larger than 256." -+ -+ if chunk_indices is None and cu_seqlens is not None: -+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) -+ if cu_seqlens is None: -+ N, NT, chunk_offsets = B, triton.cdiv(T, BT), None -+ else: -+ N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) -+ -+ dh = q.new_empty(B, NT, H, K, V) -+ dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None -+ dv2 = torch.empty_like(dv) -+ -+ BV = 128 -+ -+ g = g.permute(0, 2, 1).contiguous() -+ -+ chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[(triton.cdiv(V, BV), N * H)]( -+ q=q, -+ k=k, -+ w=w, -+ g=g, -+ gk=gk, -+ dht=dht, -+ dh0=dh0, -+ do=do, -+ dh=dh, -+ dv=dv, -+ dv2=dv2, -+ cu_seqlens=cu_seqlens, -+ chunk_offsets=chunk_offsets, -+ scale=scale, -+ T=T, -+ H=H, -+ K=K, -+ V=V, -+ BT=BT, -+ BV=BV, -+ ) -+ return dh, dh0, dv2 -diff --git a/megatron/core/ssm/triton/chunk_o.py b/megatron/core/ssm/triton/chunk_o.py -new file mode 100644 -index 000000000..9e41c2a57 ---- /dev/null -+++ b/megatron/core/ssm/triton/chunk_o.py -@@ -0,0 +1,592 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+from typing import Optional, Tuple -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices, exp, prepare_chunk_offsets -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, -+ 'USE_DW': lambda args: args['dw'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -+}) -+@triton.jit(do_not_specialize=['T']) -+def chunk_bwd_kernel_dqkwg( -+ q, -+ k, -+ v, -+ h, -+ g, -+ g_gamma, -+ do, -+ dh, -+ dq, -+ dk, -+ dg, -+ w, -+ dv, -+ dw, -+ cu_seqlens, -+ chunk_indices, -+ scale, -+ B: tl.constexpr, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ BV: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_G_GAMMA: tl.constexpr, -+ USE_DW: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ gdiff, -+): -+ i_t, i_b = tl.program_id(0), tl.program_id(1) -+ T_max = T -+ if IS_VARLEN: -+ i_tg = i_t -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ total = B * T_max -+ T = eos - bos -+ else: -+ NT = tl.cdiv(T, BT) -+ i_tg = i_b * NT + i_t -+ bos, eos = i_b * T, i_b * T + T -+ total = B * T_max -+ -+ NK = tl.cdiv(K, BK) -+ for i_k in range(NK): -+ if USE_G: -+ dg_k = dg + i_k * total * H -+ -+ for i_h in range(H): -+ v_h = v + (bos * H + i_h) * V -+ do_h = do + (bos * H + i_h) * V -+ h_h = h + (i_tg * H + i_h).to(tl.int64) * K * V -+ dh_h = dh + (i_tg * H + i_h).to(tl.int64) * K * V -+ q_h = q + (bos * H + i_h) * K -+ k_h = k + (bos * H + i_h) * K -+ dq_h = dq + (bos * H + i_h) * K -+ dk_h = dk + (bos * H + i_h) * K -+ -+ if USE_DW: -+ w_h = w + (bos * H + i_h) * K -+ dw_h = dw + (bos * H + i_h) * K -+ dv_h = dv + (bos * H + i_h) * V -+ -+ if USE_G: -+ if IS_VARLEN: -+ dg_h = dg_k + i_h * T_max + bos -+ g_h = g + i_h * T_max + bos -+ else: -+ dg_h = dg_k + (i_b * H + i_h) * T_max -+ g_h = g + (i_b * H + i_h) * T_max -+ b_dg_last = tl.zeros([1, ], dtype=tl.float32) -+ -+ if USE_G_GAMMA: -+ b_gamma = tl.load(g_gamma + i_h) -+ b_g = b_gamma * (tl.arange(0, BT) + 1) -+ b_g_last = b_gamma * min(BT, T - i_t * BT) -+ -+ b_dq = tl.zeros([BT, BK], dtype=tl.float32) -+ b_dk = tl.zeros([BT, BK], dtype=tl.float32) -+ b_ds = tl.zeros([BT, BT], dtype=tl.float32) -+ b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None -+ -+ for i_v in range(tl.cdiv(V, BV)): -+ p_v = tl.make_block_ptr(v_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_do = tl.make_block_ptr(do_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_h = tl.make_block_ptr(h_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) -+ p_dh = tl.make_block_ptr(dh_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) -+ -+ b_v = tl.load(p_v, boundary_check=(0, 1)) -+ b_do = tl.load(p_do, boundary_check=(0, 1)) -+ b_h = tl.load(p_h, boundary_check=(0, 1)) -+ b_dh = tl.load(p_dh, boundary_check=(0, 1)) -+ -+ if USE_G: -+ b_dg_last += (tl.sum(b_h * b_dh)) -+ -+ b_ds += tl.dot(b_do, tl.trans(b_v)) -+ b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) -+ b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) -+ -+ if USE_DW: -+ p_dv = tl.make_block_ptr(dv_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ b_dv = tl.load(p_dv, boundary_check=(0, 1)) -+ b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) -+ -+ if USE_DW: -+ p_dw = tl.make_block_ptr(dw_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) -+ -+ tl.debug_barrier() -+ -+ p_q = tl.make_block_ptr(q_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_k = tl.make_block_ptr(k_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ -+ p_dq = tl.make_block_ptr(dq_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_dk = tl.make_block_ptr(dk_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ -+ o_t = i_t * BT + tl.arange(0, BT) -+ m_t = o_t < T -+ m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) -+ -+ if USE_G: -+ b_dg = tl.zeros([BT, ], dtype=tl.float32) -+ p_g = tl.make_block_ptr(g_h, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_g_last = tl.load(g_h + (min(i_t * BT + BT, T) - 1) * 1) -+ b_dg_last *= tl.exp(b_g_last) -+ -+ b_dq = b_dq * tl.exp(b_g)[:, None] * scale -+ b_dg += tl.sum(b_dq * b_q, axis=1) -+ -+ b_dk = b_dk * tl.where(m_t, tl.exp(-b_g + b_g_last), 0)[:, None] -+ b_dg -= tl.sum(b_k * b_dk, axis=1) -+ b_dg_last += tl.sum(b_dk * b_k) -+ -+ if IS_VARLEN: -+ b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale -+ else: -+ p_gdiff = tl.make_block_ptr(gdiff + i_b * H * NT * BT * BT + i_h * NT * BT * BT + i_t * BT * BT, -+ (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) -+ gdiff_ = tl.load(p_gdiff) -+ b_ds = b_ds * gdiff_ * scale -+ -+ b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k)) -+ b_dg += tl.sum(b_ds2, axis=1) -+ b_dg -= tl.sum(b_ds2, axis=0) -+ -+ b_ds = b_ds.to(b_k.dtype) -+ b_dq += tl.dot(b_ds, b_k) -+ b_dk += tl.dot(tl.trans(b_ds), b_q) -+ p_dg = tl.make_block_ptr(dg_h, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ -+ last_index_local = min(BT, T - i_t * BT) - 1 -+ if last_index_local >= 0: -+ is_last_mask = tl.arange(0, BT) == last_index_local -+ b_dg = tl.where(is_last_mask, b_dg + b_dg_last, b_dg) -+ else: -+ pass -+ -+ tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) -+ tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -+ tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) -+ -+ elif USE_G_GAMMA: -+ b_dq = b_dq * exp(b_g)[:, None] * scale -+ b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] -+ b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale -+ b_ds = b_ds.to(b_k.dtype) -+ b_dq += tl.dot(b_ds, b_k) -+ b_dk += tl.dot(tl.trans(b_ds), b_q) -+ tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) -+ tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -+ -+ else: -+ b_ds = tl.where(m_A, b_ds, 0) -+ b_ds = b_ds.to(b_k.dtype) -+ b_dq += tl.dot(b_ds, b_k) -+ b_dk += tl.dot(tl.trans(b_ds), b_q) * scale -+ b_dq *= scale -+ tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) -+ tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -+}) -+@triton.jit(do_not_specialize=['T']) -+def chunk_bwd_kernel_dv_local( -+ q, -+ k, -+ g, -+ g_gamma, -+ do, -+ dv, -+ cu_seqlens, -+ chunk_indices, -+ scale, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ BV: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_G_GAMMA: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ i_t, i_b = tl.program_id(0), tl.program_id(1) -+ T_max = T -+ -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ -+ for i_h in range(H): -+ offset_kh = (bos * H + i_h) * K -+ offset_vh = (bos * H + i_h) * V -+ -+ b_A = tl.zeros([BT, BT], dtype=tl.float32) -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k + offset_kh, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_q = tl.make_block_ptr(q + offset_kh, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ b_A += tl.dot(b_k, b_q) -+ -+ if USE_G: -+ if IS_VARLEN: -+ offset_g = i_h * T_max + bos -+ else: -+ offset_g = i_b * H * T_max + i_h * T_max -+ -+ p_g = tl.make_block_ptr(g + offset_g, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ -+ if USE_G_GAMMA: -+ b_gamma = tl.load(g_gamma + i_h) -+ b_g = b_gamma * (tl.arange(0, BT) + 1) -+ -+ o_t = i_t * BT + tl.arange(0, BT) -+ m_t = o_t < T -+ m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) -+ -+ if USE_G: -+ b_A = tl.where(m_A, b_A * tl.exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) -+ else: -+ b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) -+ -+ for i_v in range(tl.cdiv(V, BV)): -+ p_do = tl.make_block_ptr(do + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_dv = tl.make_block_ptr(dv + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ b_do = tl.load(p_do, boundary_check=(0, 1)) -+ b_dv = tl.dot(b_A.to(b_do.dtype), b_do) -+ tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.jit(do_not_specialize=['T']) -+def chunk_fwd_kernel_o( -+ q, -+ k, -+ v, -+ h, -+ g, -+ g_gamma, -+ o, -+ cu_seqlens, -+ chunk_offsets, -+ scale, -+ T, -+ H: tl.constexpr, -+ N: tl.constexpr, -+ Hg: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ BV: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_G_GAMMA: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ T_max = T -+ for i_v in range(tl.cdiv(V, BV)): -+ for i_n in range(N): -+ if IS_VARLEN: -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( -+ cu_seqlens + i_n + 1 -+ ).to(tl.int32) -+ T = eos - bos -+ NT = tl.cdiv(T, BT) -+ boh = tl.load(chunk_offsets + i_n).to(tl.int64) -+ else: -+ bos, eos = i_n * T, i_n * T + T -+ NT = tl.cdiv(T, BT) -+ boh = i_n * NT -+ -+ core_id = tl.program_id(0) -+ total_cores = tl.num_programs(0) -+ base_chunks_per_pid = NT // total_cores -+ remainder = NT % total_cores -+ -+ if core_id < remainder: -+ chunks_this_pid = base_chunks_per_pid + 1 -+ start_idx = core_id * chunks_this_pid -+ else: -+ chunks_this_pid = base_chunks_per_pid -+ start_idx = core_id * base_chunks_per_pid + remainder -+ -+ # offset calculation -+ for i_h in range(0, H): -+ q_offset = (bos * Hg + i_h // (H // Hg)) * K -+ k_offset = (bos * Hg + i_h // (H // Hg)) * K -+ v_offset = (bos * H + i_h) * V -+ o_offset = (bos * H + i_h) * V -+ -+ for i_t in range(start_idx, start_idx + chunks_this_pid): -+ i_tg = boh + i_t -+ h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V -+ b_o = tl.zeros([BT, BV], dtype=tl.float32) -+ b_A = tl.zeros([BT, BT], dtype=tl.float32) -+ for i_k in range(tl.cdiv(K, BK)): -+ p_q = tl.make_block_ptr( -+ q + q_offset, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) -+ ) -+ p_k = tl.make_block_ptr( -+ k + k_offset, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) -+ ) -+ p_h = tl.make_block_ptr( -+ h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) -+ ) -+ b_q = tl.load(p_q, boundary_check=(0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ b_h = tl.load(p_h, boundary_check=(0, 1)) -+ -+ # [BT, BK] @ [BK, BV] -> [BT, BV] -+ b_o += tl.dot(b_q, b_h) -+ # [BT, BK] @ [BK, BT] -> [BT, BT] -+ b_A += tl.dot(b_q, b_k) -+ -+ if USE_G: -+ if IS_VARLEN: -+ p_g = tl.make_block_ptr(g + bos + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ else: -+ p_g = tl.make_block_ptr(g + bos * H + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_o = b_o * exp(b_g)[:, None] -+ b_A = b_A * exp(b_g[:, None] - b_g[None, :]) -+ if USE_G_GAMMA: -+ b_gamma = tl.load(g_gamma + i_h) -+ b_g = b_gamma * (tl.arange(0, BT) + 1) -+ -+ o_i = tl.arange(0, BT) -+ m_A = o_i[:, None] >= o_i[None, :] -+ b_A = tl.where(m_A, b_A, 0) -+ -+ p_v = tl.make_block_ptr( -+ v + v_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) -+ ) -+ p_o = tl.make_block_ptr( -+ o + o_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) -+ ) -+ b_v = tl.load(p_v, boundary_check=(0, 1)) -+ -+ # to fix mma -> mma layout conversion -+ # already solved by triton v3.2 or higher -+ b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale -+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+def chunk_bwd_dqkwg( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ do: torch.Tensor, -+ h: torch.Tensor, -+ dh: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ g_gamma: Optional[torch.Tensor] = None, -+ dv: Optional[torch.Tensor] = None, -+ w: Optional[torch.Tensor] = None, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64, -+ scale: float = 1.0, -+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: -+ B, T, H, K, V = *k.shape, v.shape[-1] -+ BT = min(chunk_size, max(16, triton.next_power_of_2(T))) -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ -+ BK = 128 if cu_seqlens is None else 64 -+ BV = 64 -+ NK = triton.cdiv(K, BK) -+ dq = torch.empty_like(q) -+ dk = torch.empty_like(k) -+ g = g.transpose(1, 2).contiguous() -+ dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None -+ dw = torch.empty_like(w) if w is not None else None -+ grid = (NT, B) -+ -+ if cu_seqlens is None: -+ if NT * BT == T: -+ g_ = g.reshape(B, H, NT, BT) -+ g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :] -+ g_diff = g_diff.clamp(-60, 60).exp() -+ g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device) -+ else: -+ diff = NT * BT - T -+ g_ = torch.cat((g, torch.zeros(B, H, diff).to(g.device)), dim=-1).reshape(B, H, NT, BT) -+ g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :] -+ g_diff = g_diff.clamp(-60, 60).exp() -+ g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device) -+ bias = torch.arange(0, BT).to(g.device) -+ o_t = (NT - 1) * BT + bias -+ m_t = o_t < T -+ m_A = (m_t[:, None] & m_t) -+ g_diff[:, :, -1] *= m_A -+ else: -+ g_diff = None -+ -+ chunk_bwd_kernel_dqkwg[grid]( -+ q=q, -+ k=k, -+ v=v, -+ h=h, -+ g=g, -+ g_gamma=g_gamma, -+ do=do, -+ dh=dh, -+ dv=dv, -+ w=w, -+ dw=dw, -+ dq=dq, -+ dk=dk, -+ dg=dg, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ scale=scale, -+ B=B, -+ T=T, -+ H=H, -+ K=K, -+ V=V, -+ BT=BT, -+ BK=BK, -+ BV=BV, -+ gdiff=g_diff, -+ ) -+ -+ if dg is not None: -+ dg = dg.sum(0) -+ dg = dg.transpose(1, 2).contiguous() -+ return dq, dk, dw, dg -+ -+ -+def chunk_bwd_dv_local( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ do: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ g_gamma: Optional[torch.Tensor] = None, -+ scale: float = None, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64 -+) -> torch.Tensor: -+ B, T, H, K, V = *k.shape, do.shape[-1] -+ BT = min(chunk_size, max(16, triton.next_power_of_2(T))) -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ -+ BK = 128 -+ BV = 128 -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ -+ g = g.transpose(1, 2).contiguous() -+ dv = torch.empty_like(do) -+ grid = (NT, B) -+ chunk_bwd_kernel_dv_local[grid]( -+ q=q, -+ k=k, -+ g=g, -+ g_gamma=g_gamma, -+ do=do, -+ dv=dv, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ scale=scale, -+ T=T, -+ H=H, -+ K=K, -+ V=V, -+ BT=BT, -+ BK=BK, -+ BV=BV, -+ ) -+ return dv -+ -+ -+def chunk_fwd_o( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ h: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ g_gamma: Optional[torch.Tensor] = None, -+ scale: Optional[float] = None, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64 -+) -> torch.Tensor: -+ B, T, Hg, K, V = *q.shape, v.shape[-1] -+ H = v.shape[-2] -+ BT = min(chunk_size, max(16, triton.next_power_of_2(T))) -+ chunk_indices = ( -+ prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ ) -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ if scale is None: -+ scale = k.shape[-1] ** -0.5 -+ -+ o = torch.empty_like(v) -+ if cu_seqlens is None: -+ N, chunk_offsets = B, None -+ else: -+ N, chunk_offsets = ( -+ len(cu_seqlens) - 1, -+ prepare_chunk_offsets(cu_seqlens, BT), -+ ) -+ -+ def grid(meta): -+ return (triton.cdiv(V, meta["BV"]), N * H) -+ -+ g = g.transpose(1, 2).contiguous() -+ h = h.contiguous() -+ CV_kernel_num = 24 -+ chunk_fwd_kernel_o[(CV_kernel_num,)]( -+ q, -+ k, -+ v, -+ h, -+ g, -+ g_gamma, -+ o, -+ cu_seqlens, -+ chunk_offsets, -+ scale, -+ T=T, -+ H=H, -+ N=N, -+ Hg=Hg, -+ K=K, -+ V=V, -+ BT=BT, -+ BK=128, -+ BV=128, -+ ) -+ return o -+ -+bwd_chunk_dqkwg = chunk_bwd_dqkwg -+bwd_chunk_dv_local = chunk_bwd_dv_local -diff --git a/megatron/core/ssm/triton/chunk_scaled_dot_kkt.py b/megatron/core/ssm/triton/chunk_scaled_dot_kkt.py -new file mode 100644 -index 000000000..dfbaa682d ---- /dev/null -+++ b/megatron/core/ssm/triton/chunk_scaled_dot_kkt.py -@@ -0,0 +1,337 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+from typing import Optional -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -+}) -+@triton.jit(do_not_specialize=['T', 'NT', 'TOTAL_TASKS']) -+def chunk_scaled_dot_kkt_fwd_kernel( -+ k, -+ g, -+ beta, -+ A, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ USE_G: tl.constexpr, -+ NT, -+ B, -+ TOTAL_TASKS, -+): -+ core_id = tl.program_id(0) -+ num_blocks = tl.num_programs(0) -+ T_max = T -+ -+ base_tasks_per_block = TOTAL_TASKS // num_blocks -+ remainder_tasks = TOTAL_TASKS % num_blocks -+ -+ if core_id < remainder_tasks: -+ tasks_this_core = base_tasks_per_block + 1 -+ start_idx = core_id * tasks_this_core -+ else: -+ tasks_this_core = base_tasks_per_block -+ start_idx = core_id * base_tasks_per_block + remainder_tasks -+ -+ for idx in range(start_idx, start_idx + tasks_this_core): -+ i_b = idx // NT -+ local_idx = idx % NT -+ -+ if IS_VARLEN: -+ i_n = tl.load(chunk_indices + local_idx * 2).to(tl.int32) -+ i_t = tl.load(chunk_indices + local_idx * 2 + 1).to(tl.int32) -+ bos = tl.load(cu_seqlens + i_n).to(tl.int32) -+ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T_local = eos - bos -+ else: -+ bos, eos = 0, T -+ i_t = local_idx -+ T_local = T -+ -+ for i_h in range(H): -+ k_batch_off = i_b * T_max * H * K -+ beta_batch_off = i_b * H * T_max -+ g_batch_off = i_b * H * T_max -+ A_batch_off = i_b * T_max * H * BT -+ -+ p_beta = tl.make_block_ptr(beta + beta_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_beta = tl.load(p_beta, boundary_check=(0,)) -+ -+ b_A = tl.zeros([BT, BT], dtype=tl.float32) -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k + k_batch_off + (bos * H + i_h) * K, (T_local, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ dot_product = tl.dot(b_k, tl.trans(b_k)) -+ -+ o_t = i_t * BT + tl.arange(0, BT) -+ o_t = o_t.to(tl.float32) -+ T_mask = (o_t < T_local).to(tl.float32) -+ -+ row_indices = tl.arange(0, BT)[:, None] -+ col_indices = tl.arange(0, BT)[None, :] -+ tril_mask = (row_indices > col_indices).to(tl.float32) -+ tril_mask = tril_mask * T_mask[:, None] -+ masked_dot = dot_product * tril_mask -+ b_A += masked_dot -+ -+ if USE_G: -+ p_g = tl.make_block_ptr(g + g_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_g_diff = b_g[:, None] - b_g[None, :] -+ b_g_diff = tl.minimum(tl.maximum(b_g_diff, -50.0), 50.0) -+ b_A *= tl.exp(b_g_diff) -+ b_A *= b_beta[:, None] -+ -+ p_A = tl.make_block_ptr(A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) -+ tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+@triton.heuristics({ -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.autotune( -+ configs=[ -+ triton.Config({'BK': BK}) -+ for BK in [32, 64] -+ ], -+ key=["BC"] -+) -+@triton.jit(do_not_specialize=['T']) -+def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( -+ k, -+ g, -+ beta, -+ A, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ BT: tl.constexpr, -+ BC: tl.constexpr, -+ BK: tl.constexpr, -+ NC: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ i_t, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) -+ i_i, i_j = i_c // NC, i_c % NC -+ -+ for i_h in range(H): -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T_val = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ T_val = T -+ -+ should_compute = (i_t * BT + i_i * BC < T_val) and (i_i > i_j) -+ -+ if should_compute: -+ k_ptr = k + (bos * H + i_h) * K -+ g_ptr = g + (bos * H + i_h) * K -+ A_ptr = A + (bos * H + i_h) * BT -+ -+ p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T_val,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) -+ b_beta = tl.load(p_beta, boundary_check=(0,)) -+ -+ b_A = tl.zeros([BC, BC], dtype=tl.float32) -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), -+ (1, 0)) -+ p_g = tl.make_block_ptr(g_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), -+ (1, 0)) -+ b_kt = tl.make_block_ptr(k_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), -+ (0, 1)) -+ p_gk = tl.make_block_ptr(g_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), -+ (0, 1)) -+ -+ o_k = i_k * BK + tl.arange(0, BK) -+ m_k = o_k < K -+ b_gn = tl.load(g_ptr + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) -+ b_g = tl.load(p_g, boundary_check=(0, 1)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.exp(b_g - b_gn[None, :]) -+ b_gk = tl.load(p_gk, boundary_check=(0, 1)) -+ b_kt = tl.load(b_kt, boundary_check=(0, 1)) * tl.exp(b_gn[:, None] - b_gk) -+ b_A += tl.dot(b_k, b_kt) -+ b_A *= b_beta[:, None] -+ -+ p_A = tl.make_block_ptr(A_ptr, (T_val, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) -+ tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+@triton.heuristics({ -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.jit(do_not_specialize=['T']) -+def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( -+ k, -+ g, -+ beta, -+ A, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ BT: tl.constexpr, -+ BC: tl.constexpr, -+ BK: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ i_t, i_i, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) -+ -+ for i_h in range(H): -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T_val = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ T_val = T -+ -+ should_compute = (i_t * BT + i_i * BC < T_val) -+ -+ if should_compute: -+ o_i = tl.arange(0, BC) -+ o_k = tl.arange(0, BK) -+ m_k = o_k < K -+ m_A = (i_t * BT + i_i * BC + o_i) < T_val -+ o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC -+ -+ p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), -+ (1, 0)) -+ p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), -+ (1, 0)) -+ p_beta = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h -+ -+ b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_beta, mask=m_A, other=0)[:, None] -+ b_g = tl.load(p_g, boundary_check=(0, 1)) -+ -+ p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k -+ p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k -+ -+ for j in range(0, min(BC, T_val - i_t * BT - i_i * BC)): -+ b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) -+ b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) -+ b_A = tl.sum(b_k * b_kt[None, :] * tl.exp(b_g - b_gk[None, :]), 1) -+ # 转化成f32 -+ o_i_tmp = o_i.to(tl.float32) -+ b_A = tl.where(o_i_tmp > j, b_A, 0.) -+ -+ tl.store(A + o_A + j, b_A, mask=m_A) -+ p_kt += H * K -+ p_gk += H * K -+ -+ -+def chunk_scaled_dot_kkt_fwd( -+ k: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ gk: Optional[torch.Tensor] = None, -+ beta: Optional[torch.Tensor] = None, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+ chunk_size: int = 64, -+ output_dtype: torch.dtype = torch.float32 -+) -> torch.Tensor: -+ r""" -+ Compute beta * K * K^T. -+ -+ Args: -+ k (torch.Tensor): -+ The key tensor of shape `[B, T, H, K]`. -+ beta (torch.Tensor): -+ The beta tensor of shape `[B, T, H]`. -+ g (torch.Tensor): -+ The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. -+ gk (torch.Tensor): -+ The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. -+ cu_seqlens (torch.LongTensor): -+ The cumulative sequence lengths of the input tensor. -+ Default: None -+ chunk_size (int): -+ The chunk size. Default: 64. -+ output_dtype (torch.dtype): -+ The dtype of the output tensor. Default: `torch.float32` -+ -+ Returns: -+ beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. -+ """ -+ B, T, H, K = k.shape -+ BT = chunk_size -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ beta = beta.transpose(1, 2).contiguous() -+ g = g.transpose(1, 2).contiguous() -+ BK = 128 -+ kernel_num = 24 -+ -+ if gk is None: -+ A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) -+ chunk_scaled_dot_kkt_fwd_kernel[(kernel_num,)]( -+ k=k, -+ g=g, -+ beta=beta, -+ A=A, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ K=K, -+ BT=BT, -+ BK=BK, -+ NT=NT, -+ B=B, -+ TOTAL_TASKS=B * NT, -+ ) -+ return A -+ -+ BC = min(16, BT) -+ NC = triton.cdiv(BT, BC) -+ BK = max(triton.next_power_of_2(K), 16) -+ A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) -+ grid = (NT, NC * NC, B) -+ chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( -+ k=k, -+ g=gk, -+ beta=beta, -+ A=A, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ K=K, -+ BT=BT, -+ BC=BC, -+ NC=NC, -+ ) -+ -+ grid = (NT, NC, B) -+ chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( -+ k=k, -+ g=gk, -+ beta=beta, -+ A=A, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ K=K, -+ BT=BT, -+ BC=BC, -+ BK=BK, -+ ) -+ return A -diff --git a/megatron/core/ssm/triton/convolution.py b/megatron/core/ssm/triton/convolution.py -new file mode 100644 -index 000000000..30eba6c92 ---- /dev/null -+++ b/megatron/core/ssm/triton/convolution.py -@@ -0,0 +1,980 @@ -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Wenshuo Zhao -+# Copyright (c) 2026, Huawei Technologies Co., Ltd. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+ -+# pylint: disable=no-name-in-module,relative-beyond-top-level -+ -+from typing import Optional -+ -+import torch -+import triton -+import triton.language as tl -+ -+ -+# Compatibility shim: in newer triton-ascend, slice ops live in -+# `triton.language.extra.cann.extension` instead of `triton.language`. -+# Re-expose them on `tl` so kernel code works on both versions. -+try: -+ from triton.language.extra.cann.extension import extract_slice, insert_slice -+ -+ if not hasattr(tl, "extract_slice"): -+ tl.extract_slice = extract_slice -+ if not hasattr(tl, "insert_slice"): -+ tl.insert_slice = insert_slice -+except ImportError: -+ pass -+ -+from .utils import get_vector_num, input_guard, prepare_chunk_indices -+ -+ -+@triton.heuristics( -+ { -+ "HAS_WEIGHT": lambda args: args["weight"] is not None, -+ "HAS_BIAS": lambda args: args["bias"] is not None, -+ "HAS_RESIDUAL": lambda args: args["residual"] is not None, -+ "USE_INITIAL_STATE": lambda args: args["initial_state"] is not None, -+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, -+ } -+) -+@triton.jit(do_not_specialize=['T', 'NUM_CHKS']) -+def causal_conv1d_fwd_kernel( -+ x, -+ y, -+ weight, -+ bias, -+ residual, -+ cu_seqlens, -+ initial_state, -+ chunk_indices, -+ B, -+ T, -+ D: tl.constexpr, -+ W: tl.constexpr, -+ BT: tl.constexpr, -+ BD: tl.constexpr, -+ ACTIVATION: tl.constexpr, -+ HAS_WEIGHT: tl.constexpr, -+ HAS_BIAS: tl.constexpr, -+ HAS_RESIDUAL: tl.constexpr, -+ USE_INITIAL_STATE: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ NUM_CHKS: tl.int32, -+ NUM_BLKS_D: tl.int32, -+): -+ pid = tl.program_id(0) -+ num_programs = tl.num_programs(0) -+ -+ total_tasks = NUM_BLKS_D * NUM_CHKS -+ -+ for task_id in range(pid, total_tasks, num_programs): -+ i_d_blk = task_id % NUM_BLKS_D -+ i_chk = task_id // NUM_BLKS_D -+ -+ i_d = i_d_blk -+ -+ if IS_VARLEN: -+ idx_ptr = chunk_indices + i_chk * 2 -+ i_n = tl.load(idx_ptr).to(tl.int32) -+ i_t = tl.load(idx_ptr + 1).to(tl.int32) -+ -+ bos = tl.load(cu_seqlens + i_n).to(tl.int64) -+ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) -+ T_len = eos - bos -+ else: -+ NT_per_seq = tl.cdiv(T, BT) -+ i_b = i_chk // NT_per_seq -+ i_t = i_chk % NT_per_seq -+ -+ i_n = i_b -+ bos = (i_b * T).to(tl.int64) -+ eos = (i_b * T + T).to(tl.int64) -+ T_len = T -+ -+ o_d = i_d * BD + tl.arange(0, BD) -+ m_d = o_d < D -+ -+ # Tail-of-allocation guard: block end in absolute packed rows must not -+ # exceed B*T, else MTE DMA touches unmapped pages. -+ is_tail_chunk = (bos + i_t * BT + BT) > (B * T) -+ -+ if HAS_WEIGHT: -+ p_w = tl.make_block_ptr(weight, (W, D), (D, 1), (0, i_d * BD), (W, BD), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1)) -+ -+ b_y = tl.zeros((BT, BD), dtype=tl.float32) -+ -+ yi_offset_1 = i_d * BD + tl.arange(0, BD)[None, :] -+ -+ if not USE_INITIAL_STATE: -+ for i_w in tl.static_range(-W + 1, 1): -+ yi_offset_0 = i_t * BT + i_w + tl.arange(0, BT)[:, None] -+ -+ mask = (yi_offset_0 < T_len) & (yi_offset_1 < D) & (yi_offset_0 >= 0) -+ # We keep intra loop load because preloading will cause ub overflow under certain tiling. -+ b_yi = tl.load(x + bos * D + yi_offset_0 * D + yi_offset_1, mask=mask, other=0.0).to(tl.float32) -+ if HAS_WEIGHT: -+ b_yi *= tl.extract_slice(b_w, [i_w + W - 1, 0], [1, BD], [1, 1]) -+ -+ b_y += b_yi -+ elif i_t * BT >= W: -+ for i_w in tl.static_range(-W + 1, 1): -+ yi_offset_0 = i_t * BT + i_w + tl.arange(0, BT)[:, None] -+ mask = (yi_offset_0 < T_len) & (yi_offset_1 < D) & (yi_offset_0 >= 0) -+ b_yi = tl.load(x + bos * D + yi_offset_0 * D + yi_offset_1, mask=mask, other=0.0).to(tl.float32) -+ if HAS_WEIGHT: -+ b_yi *= tl.extract_slice(b_w, [i_w + W - 1, 0], [1, BD], [1, 1]) -+ b_y += b_yi -+ else: -+ o_t = i_t * BT + tl.arange(0, BT) -+ for i_w in tl.static_range(-W + 1, 1): -+ o_x = o_t + i_w -+ -+ m_x = ((o_x >= 0) & (o_x < T_len))[:, None] & m_d -+ -+ m_c = ((o_x + W >= 0) & (o_x < 0))[:, None] & m_d -+ -+ b_yi = tl.load(x + bos * D + o_x[:, None] * D + o_d, mask=m_x, other=0).to(tl.float32) -+ -+ b_yi += tl.load(initial_state + i_n * D * W + o_d * W + (o_x + W)[:, None], mask=m_c, other=0).to( -+ tl.float32 -+ ) -+ -+ if HAS_WEIGHT: -+ b_yi *= tl.extract_slice(b_w, [i_w + W - 1, 0], [1, BD], [1, 1]) -+ b_y += b_yi -+ -+ if HAS_BIAS: -+ b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32) -+ -+ if ACTIVATION == 'swish' or ACTIVATION == 'silu': # pylint: disable=consider-using-in -+ b_y = b_y * tl.sigmoid(b_y) -+ -+ if HAS_RESIDUAL: -+ if is_tail_chunk: -+ o_t_r = i_t * BT + tl.arange(0, BT) -+ m_t_r = (o_t_r >= 0) & (o_t_r < T_len) -+ b_residual = tl.load( -+ residual + bos * D + o_t_r[:, None] * D + o_d[None, :], -+ mask=m_t_r[:, None] & m_d[None, :], -+ other=0.0, -+ ) -+ else: -+ p_residual = tl.make_block_ptr( -+ residual + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0) -+ ) -+ b_residual = tl.load(p_residual, boundary_check=(0, 1)) -+ b_y += b_residual -+ -+ if is_tail_chunk: -+ o_t_y = i_t * BT + tl.arange(0, BT) -+ m_t_y = (o_t_y >= 0) & (o_t_y < T_len) -+ b_y_cast = tl.cast(b_y, dtype=y.dtype.element_ty, fp_downcast_rounding="rtne") -+ tl.store( -+ y + bos * D + o_t_y[:, None] * D + o_d[None, :], -+ b_y_cast, -+ mask=m_t_y[:, None] & m_d[None, :], -+ ) -+ else: -+ p_y = tl.make_block_ptr(y + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) -+ tl.store(p_y, tl.cast(b_y, dtype=p_y.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) -+ -+ -+@triton.heuristics( -+ { -+ "HAS_WEIGHT": lambda args: args["dw"] is not None, -+ "HAS_BIAS": lambda args: args["db"] is not None, -+ "USE_INITIAL_STATE": lambda args: args["dh0"] is not None, -+ "USE_FINAL_STATE": lambda args: args["dht"] is not None, -+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, -+ } -+) -+@triton.jit(do_not_specialize=['T', 'NUM_CHKS']) -+def causal_conv1d_bwd_kernel( -+ x, -+ y, -+ weight, -+ initial_state, -+ dh0, -+ dht, -+ dy, -+ dx, -+ dw, -+ db, -+ cu_seqlens, -+ chunk_indices, -+ B, -+ T, -+ D: tl.constexpr, -+ W: tl.constexpr, -+ BT: tl.constexpr, -+ BD: tl.constexpr, -+ ACTIVATION: tl.constexpr, -+ HAS_WEIGHT: tl.constexpr, -+ HAS_BIAS: tl.constexpr, -+ USE_INITIAL_STATE: tl.constexpr, -+ USE_FINAL_STATE: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ NUM_BLKS_D: tl.int32, -+ NUM_CHKS: tl.int32, -+): -+ pid = tl.program_id(0) -+ num_programs = tl.num_programs(0) -+ -+ # Total packed rows = allocation upper bound of x / dy / dx, used to -+ # detect tail chunks whose block end would overshoot the packed tensor -+ # and trigger MTE "DDR address out of range". varlen or not, the tensor -+ # shape is always [B, T, D]. -+ TOTAL_ROWS = B * T -+ -+ total_tasks = NUM_CHKS * NUM_BLKS_D -+ -+ for task_id in range(pid, total_tasks, num_programs): # pylint: disable=too-many-nested-blocks -+ i_d = task_id % NUM_BLKS_D -+ i_chk = task_id // NUM_BLKS_D -+ -+ if IS_VARLEN: -+ i_t = i_chk -+ -+ idx_chk = i_chk -+ -+ i_tg = idx_chk -+ -+ ptr = chunk_indices + idx_chk * 2 -+ i_n = tl.load(ptr).to(tl.int32) -+ i_t_offset = tl.load(ptr + 1).to(tl.int32) -+ -+ i_t = i_t_offset -+ -+ bos = tl.load(cu_seqlens + i_n).to(tl.int64) -+ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) -+ T_len = eos - bos -+ else: -+ NT_per_seq = tl.cdiv(T, BT) -+ -+ i_b = i_chk // NT_per_seq -+ i_t = i_chk % NT_per_seq -+ -+ i_tg = i_chk -+ -+ i_n = i_b -+ bos = (i_b * T).to(tl.int64) -+ eos = (i_b * T + T).to(tl.int64) -+ T_len = T -+ -+ o_d = i_d * BD + tl.arange(0, BD) -+ m_d = o_d < D -+ -+ is_tail_chunk = (bos + i_t * BT + BT * W) > TOTAL_ROWS -+ -+ if HAS_WEIGHT: -+ if is_tail_chunk: -+ o_t_x = i_t * BT + tl.arange(0, BT) -+ m_t_x = (o_t_x >= 0) & (o_t_x < T_len) -+ b_x = tl.load( -+ x + bos * D + o_t_x[:, None] * D + o_d[None, :], -+ mask=m_t_x[:, None] & m_d[None, :], -+ other=0, -+ ) -+ else: -+ p_x = tl.make_block_ptr(x + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) -+ b_x = tl.load(p_x, boundary_check=(0, 1)) -+ -+ p_w = tl.make_block_ptr(weight, (W, D), (D, 1), (0, i_d * BD), (W, BD), (1, 0)) -+ b_w = tl.load(p_w, boundary_check=(0, 1), padding_option="zero") -+ -+ b_dx = tl.zeros((BT, BD), dtype=tl.float32) -+ if HAS_BIAS: -+ b_db = tl.zeros((BD,), dtype=tl.float32) -+ -+ if not USE_FINAL_STATE and not USE_INITIAL_STATE: -+ b_dw = tl.zeros((W, BD), dtype=tl.float32) -+ -+ if is_tail_chunk: -+ o_t_full = i_t * BT + tl.arange(0, BT * W) -+ m_t_full = (o_t_full >= 0) & (o_t_full < T_len) -+ b_dy = tl.load( -+ dy + bos * D + o_t_full[:, None] * D + o_d[None, :], -+ mask=m_t_full[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ b_y = tl.load( -+ y + bos * D + o_t_full[:, None] * D + o_d[None, :], -+ mask=m_t_full[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ else: -+ p_dy = tl.make_block_ptr(dy + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT * W, BD), (1, 0)) -+ b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) -+ -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ p_y = tl.make_block_ptr(y + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT * W, BD), (1, 0)) -+ b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) -+ -+ for i_w in tl.static_range(0, W): -+ b_dy_sub = tl.extract_slice(b_dy, [i_w, 0], [BT, BD], [1, 1]) -+ -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ b_y_sub = tl.extract_slice(b_y, [i_w, 0], [BT, BD], [1, 1]) # pylint: disable=used-before-assignment -+ b_ys = tl.sigmoid(b_y_sub) -+ b_dy_sub = b_dy_sub * b_ys * (1 + b_y_sub * (1 - b_ys)) -+ -+ b_wdy = b_dy_sub -+ if HAS_WEIGHT: -+ b_wdy = b_wdy * tl.extract_slice(b_w, [W - i_w - 1, 0], [1, BD], [1, 1]) -+ -+ b_dw_sub = tl.sum(b_dy_sub * b_x, 0) # [BT, BD] * [BT, BD] --> sum(0) = [BD] -+ b_dw = tl.insert_slice(b_dw, b_dw_sub[None, :], [W - i_w - 1, 0], [1, BD], [1, 1]) -+ -+ if HAS_BIAS and i_w == 0: -+ b_db += tl.sum(b_dy_sub, 0) -+ b_dx += b_wdy -+ -+ p_dw = tl.make_block_ptr(dw + i_tg * W * D, (W, D), (D, 1), (0, i_d * BD), (W, BD), (1, 0)) -+ tl.store(p_dw, b_dw.to(dw.dtype.element_ty)) -+ elif i_t * BT >= W: -+ for i_w in tl.static_range(0, W): -+ if is_tail_chunk: -+ o_t_iw = i_t * BT + i_w + tl.arange(0, BT) -+ m_t_iw = (o_t_iw >= 0) & (o_t_iw < T_len) -+ b_dy = tl.load( -+ dy + bos * D + o_t_iw[:, None] * D + o_d[None, :], -+ mask=m_t_iw[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ b_y = tl.load( -+ y + bos * D + o_t_iw[:, None] * D + o_d[None, :], -+ mask=m_t_iw[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ b_ys = tl.sigmoid(b_y) -+ b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys)) -+ else: -+ p_dy = tl.make_block_ptr( -+ dy + bos * D, (T_len, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0) -+ ) -+ b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ p_y = tl.make_block_ptr( -+ y + bos * D, (T_len, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0) -+ ) -+ b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) -+ b_ys = tl.sigmoid(b_y) -+ b_dy = b_dy * b_ys * (1 + b_y * (1 - b_ys)) -+ b_wdy = b_dy -+ if HAS_WEIGHT: -+ b_wdy = b_wdy * tl.extract_slice(b_w, [W - i_w - 1, 0], [1, BD], [1, 1]) -+ -+ b_dw = tl.sum(b_dy * b_x, 0) -+ tl.store(dw + i_tg * W * D + (W - i_w - 1) * D + o_d, b_dw.to(dw.dtype.element_ty), mask=m_d) -+ if HAS_BIAS and i_w == 0: -+ b_db += tl.sum(b_dy, 0) -+ b_dx += b_wdy -+ else: -+ o_t = i_t * BT + tl.arange(0, BT) -+ for i_w in tl.static_range(0, W): -+ if is_tail_chunk: -+ o_t_iw = i_t * BT + i_w + tl.arange(0, BT) -+ m_t_iw = (o_t_iw >= 0) & (o_t_iw < T_len) -+ b_dy_shift = tl.load( -+ dy + bos * D + o_t_iw[:, None] * D + o_d[None, :], -+ mask=m_t_iw[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ b_y = tl.load( -+ y + bos * D + o_t_iw[:, None] * D + o_d[None, :], -+ mask=m_t_iw[:, None] & m_d[None, :], -+ other=0.0, -+ ).to(tl.float32) -+ b_ys = tl.sigmoid(b_y) -+ b_dy_shift = b_dy_shift * b_ys * (1 + b_y * (1 - b_ys)) -+ else: -+ p_dy = tl.make_block_ptr( -+ dy + bos * D, (T_len, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0) -+ ) -+ b_dy_shift = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ p_y = tl.make_block_ptr( -+ y + bos * D, (T_len, D), (D, 1), (i_t * BT + i_w, i_d * BD), (BT, BD), (1, 0) -+ ) -+ b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) -+ b_ys = tl.sigmoid(b_y) -+ b_dy_shift = b_dy_shift * b_ys * (1 + b_y * (1 - b_ys)) -+ if HAS_WEIGHT: -+ b_dw = tl.sum(b_dy_shift * b_x, 0) -+ -+ if USE_INITIAL_STATE: -+ mask_head_rows = o_t < i_w -+ -+ b_dy_head = tl.load( -+ dy + bos * D + o_t[:, None] * D + o_d, -+ mask=(mask_head_rows[:, None] & m_d[None, :]), -+ other=0.0, -+ ).to(tl.float32) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ b_y_head = tl.load( -+ y + bos * D + o_t[:, None] * D + o_d, -+ mask=(mask_head_rows[:, None] & m_d[None, :]), -+ other=0.0, -+ ).to(tl.float32) -+ b_ys_head = tl.sigmoid(b_y_head) -+ b_dy_head = b_dy_head * b_ys_head * (1 + b_y_head * (1 - b_ys_head)) -+ o_c = W - i_w + o_t -+ -+ mask_c = mask_head_rows & (o_c >= 1) & (o_c < W) -+ b_xc = tl.load( -+ initial_state + i_n * D * W + o_d[None, :] * W + o_c[:, None], -+ mask=(mask_c[:, None] & m_d[None, :]), -+ other=0.0, -+ ).to(tl.float32) -+ -+ b_dw += tl.sum(b_dy_head * b_xc, 0) -+ tl.store(dw + i_tg * W * D + (W - i_w - 1) * D + o_d, b_dw.to(dw.dtype.element_ty), mask=m_d) -+ -+ if HAS_BIAS and i_w == 0: -+ b_db += tl.sum(b_dy_shift, 0) -+ b_wdy = ( -+ b_dy_shift -+ if not HAS_WEIGHT -+ else (b_dy_shift * tl.extract_slice(b_w, [W - i_w - 1, 0], [1, BD], [1, 1])) -+ ) -+ b_dx += b_wdy -+ -+ if USE_INITIAL_STATE: -+ for i_w in tl.static_range(1, W): -+ # dh0[i_w] = sum_{t=0}^{i_w-1} dy0[t, :] * w[i_w-1-t, :] -+ # 逐行 load dy0 避免预加载 [BT,BD] 炸 UB，消除三维 i1 broadcast -+ b_dh0_s = tl.zeros((BD,), dtype=tl.float32) -+ for i_t2 in tl.static_range(0, W - 1): -+ if i_t2 < i_w: -+ dy0_row = tl.load(dy + bos * D + (i_t * BT + i_t2) * D + o_d, mask=m_d, other=0.0).to( -+ tl.float32 -+ ) -+ if ACTIVATION == "swish" or ACTIVATION == "silu": # pylint: disable=consider-using-in -+ y0_row = tl.load(y + bos * D + (i_t * BT + i_t2) * D + o_d, mask=m_d, other=0.0).to( -+ tl.float32 -+ ) -+ y0_s = tl.sigmoid(y0_row) -+ dy0_row = dy0_row * y0_s * (1 + y0_row * (1 - y0_s)) -+ if HAS_WEIGHT: -+ w_row = tl.extract_slice(b_w, [i_w - 1 - i_t2, 0], [1, BD], [1, 1]) -+ b_dh0_s += tl.sum(dy0_row[None, :] * w_row, 0).to(tl.float32) -+ else: -+ b_dh0_s += dy0_row -+ -+ tl.store( -+ dh0 + i_t * B * D * W + i_n * D * W + o_d * W + i_w, -+ b_dh0_s.to(dh0.dtype.element_ty, fp_downcast_rounding="rtne"), -+ mask=m_d, -+ ) -+ -+ if HAS_BIAS: -+ b_db = tl.cast(b_db, dtype=db.dtype.element_ty, fp_downcast_rounding="rtne") -+ tl.store(db + i_tg * D + o_d, b_db, mask=m_d) -+ -+ if USE_FINAL_STATE: -+ if i_t * BT + BT >= T_len - W: -+ # final_state[b,d,w] = x[b, T_len-W+w, d]，w ∈ [0, W) -+ # 所以 dx[t] += dht[b, d, t-(T_len-W)]，当 t ∈ [T_len-W, T_len-1] -+ row_arange = tl.arange(0, BT) -+ for i_w in tl.static_range(0, W): -+ target_row = T_len - W + i_w -+ local_row = target_row - i_t * BT -+ in_chunk = (local_row >= 0) & (local_row < BT) & (target_row >= 0) & (target_row < T_len) -+ b_dht_row = tl.load( -+ dht + i_n * D * W + o_d * W + i_w, -+ mask=m_d, -+ other=0.0, -+ ).to(tl.float32) -+ row_match = (row_arange == local_row) & in_chunk -+ b_dx += tl.where( -+ row_match[:, None] & m_d[None, :], -+ b_dht_row[None, :], -+ 0.0, -+ ) -+ -+ if is_tail_chunk: -+ o_t_dx = i_t * BT + tl.arange(0, BT) -+ m_t_dx = (o_t_dx >= 0) & (o_t_dx < T_len) -+ b_dx_cast = tl.cast(b_dx, dtype=dx.dtype.element_ty, fp_downcast_rounding="rtne") -+ tl.store( -+ dx + bos * D + o_t_dx[:, None] * D + o_d[None, :], -+ b_dx_cast, -+ mask=m_t_dx[:, None] & m_d[None, :], -+ ) -+ else: -+ p_dx = tl.make_block_ptr(dx + bos * D, (T_len, D), (D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) -+ tl.store( -+ p_dx, tl.cast(b_dx, dtype=p_dx.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1) -+ ) -+ -+ -+@input_guard -+def causal_conv1d_fwd_impl( -+ x: torch.Tensor, -+ weight: torch.Tensor, -+ bias: torch.Tensor, -+ residual: torch.Tensor, -+ initial_state: Optional[torch.Tensor] = None, -+ output_final_state: bool = False, -+ activation: Optional[str] = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+) -> torch.Tensor: -+ shape = x.shape -+ if x.shape[-1] != weight.shape[-1]: -+ raise ValueError("x [B, T, D], weight [W, D], please check.") -+ B, T, D, W = *x.shape, weight.shape[0] -+ NUM_CORES = get_vector_num() -+ # USE_INITIAL_STATE: the else-branch (first chunk) uses tl.static_range which -+ # unrolls W iterations, each keeping both x-load and initial_state-load live. -+ # Combined with NPU multi-buffering this easily overflows the ~192 KB UB. -+ # Reduce BD and cap BT to keep peak UB within budget. -+ if initial_state is not None: -+ BD = 32 -+ BT = min(16, triton.next_power_of_2(triton.cdiv(max(16, B * T), NUM_CORES))) -+ else: -+ BD = 256 -+ BT = min(32, triton.next_power_of_2(triton.cdiv(max(16, B * T), NUM_CORES))) -+ if D % BD != 0: -+ raise ValueError("D must be divisible by BD.") -+ NUM_BLKS_D = triton.cdiv(D, BD) -+ -+ if cu_seqlens is not None: -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) -+ NUM_CHKS = len(chunk_indices) -+ else: -+ chunk_indices = None -+ -+ NUM_CHKS = triton.cdiv(T, BT) * B -+ -+ y = torch.empty_like(x) -+ -+ grid = (NUM_CORES,) -+ -+ causal_conv1d_fwd_kernel[grid]( -+ x=x, -+ y=y, -+ weight=weight, -+ bias=bias, -+ residual=residual, -+ cu_seqlens=cu_seqlens, -+ initial_state=initial_state, -+ chunk_indices=chunk_indices, -+ B=B, -+ T=T, -+ D=D, -+ W=W, -+ BT=BT, -+ BD=BD, -+ ACTIVATION=activation, -+ NUM_CHKS=NUM_CHKS, -+ NUM_BLKS_D=NUM_BLKS_D, -+ ) -+ -+ final_state = None -+ if output_final_state: -+ final_state = causal_conv1d_update_states( -+ x=x, -+ state_len=W, -+ initial_state=initial_state, -+ cu_seqlens=cu_seqlens, -+ ) -+ -+ return y.view(shape), final_state -+ -+ -+@input_guard -+def causal_conv1d_bwd_impl( -+ x: torch.Tensor, -+ dy: torch.Tensor, -+ dht: torch.Tensor, -+ weight: Optional[torch.Tensor] = None, -+ bias: Optional[torch.Tensor] = None, -+ residual: Optional[torch.Tensor] = None, -+ initial_state: Optional[torch.Tensor] = None, -+ activation: str = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+): -+ shape = x.shape -+ if x.shape[-1] != weight.shape[-1]: -+ raise ValueError("x [B, T, D], weight [W, D], please check.") -+ -+ B, T, D = x.shape -+ W = weight.shape[0] if weight is not None else None -+ -+ NUM_CORES = get_vector_num() -+ # ---- UB-aware tile sizing for backward ---- -+ # UB capacity: 192 KB = 196608 bytes. With multi-buffering (up to 3x), -+ # effective budget ≈ 64 KB per "live set". -+ # -+ # Path C (USE_INITIAL_STATE, worst case with activation) peak live buffers: -+ # b_x[BT,BD], b_w[W,BD] — input dtype (es bytes each) -+ # b_dx[BT,BD], b_dy_shift[BT,BD], b_y[BT,BD], -+ # b_dy_head[BT,BD], b_xc[BT,BD] — fp32 (4 bytes each) -+ # b_dw[BD], b_db[BD] — fp32 (small) -+ # -+ # Peak ≈ BT*BD*(2*es + 5*4) + W*BD*es + 2*BD*4 -+ # ≈ BT*BD*(es*2 + 20) + W*BD*es (ignoring small terms) -+ # -+ # Budget: BT*BD*(es*2 + 20) + W*BD*es ≤ 65536 bytes -+ # -+ # With BT=8: BD ≤ 65536 / (8*(4*2+20) + W*4) = 65536 / (8*28 + W*4) -+ # W=4: BD ≤ 65536 / 240 = 273 → clamp to 256 (but too aggressive) -+ # Conservative: fix BT=8, BD=32 gives 8*32*28 + 4*32*4 = 7168+512 = 7.5KB ✓✓✓ -+ # -+ # We use BT=8, BD=32 as safe defaults that work for all W≤8 and all dtypes. -+ # This matches the front-end's conservative approach but accounts for the -+ # extra gradient buffers in backward. -+ if initial_state is not None: -+ BD = 32 -+ BT = min(8, triton.next_power_of_2(triton.cdiv(max(16, B * T), NUM_CORES))) -+ else: -+ BD = 32 -+ BT = min(32, triton.next_power_of_2(triton.cdiv(max(16, B * T), NUM_CORES))) -+ if D % BD != 0: -+ raise ValueError("D must be divisible by BD.") -+ NUM_BLKS_D = triton.cdiv(D, BD) -+ -+ if cu_seqlens is not None: -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) -+ NUM_CHKS = len(chunk_indices) -+ -+ NT = len(chunk_indices) -+ else: -+ chunk_indices = None -+ -+ NT = triton.cdiv(T, BT) -+ NUM_CHKS = NT * B -+ -+ y = None -+ if activation is not None: -+ y, _ = causal_conv1d_fwd_impl( -+ x=x, -+ weight=weight, -+ bias=bias, -+ residual=None, -+ initial_state=initial_state, -+ activation=None, -+ cu_seqlens=cu_seqlens, -+ output_final_state=False, -+ ) -+ dx = torch.empty_like(x) -+ dw = weight.new_empty(B * NT, W, D, dtype=torch.float) if weight is not None else None -+ db = bias.new_empty(B * NT, *bias.shape, dtype=torch.float) if bias is not None else None -+ dr = dy if residual is not None else None -+ -+ if initial_state is not None: -+ if cu_seqlens is not None: -+ eff_NT = len(chunk_indices) -+ else: -+ eff_NT = triton.cdiv(T, BT) -+ -+ dh0 = initial_state.new_zeros(min(eff_NT, triton.cdiv(W, BT)), *initial_state.shape) -+ else: -+ dh0 = None -+ -+ grid = (NUM_CORES,) -+ -+ causal_conv1d_bwd_kernel[grid]( -+ x=x, -+ y=y, -+ weight=weight, -+ initial_state=initial_state, -+ dh0=dh0, -+ dht=dht, -+ dy=dy, -+ dx=dx, -+ dw=dw, -+ db=db, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ B=B, -+ T=T, -+ D=D, -+ W=W, -+ BT=BT, -+ BD=BD, -+ ACTIVATION=activation, -+ NUM_BLKS_D=NUM_BLKS_D, -+ NUM_CHKS=NUM_CHKS, -+ ) -+ -+ if weight is not None: -+ dw = dw.sum(0).contiguous().to(weight) -+ if bias is not None: -+ db = db.sum(0).to(bias) -+ if initial_state is not None: -+ dh0 = dh0.sum(0, dtype=torch.float32).to(initial_state) -+ -+ return dx.view(shape), dw, db, dr, dh0 -+ -+ -+@triton.heuristics( -+ { -+ "USE_INITIAL_STATE": lambda args: args["initial_state"] is not None, -+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, -+ } -+) -+@triton.jit(do_not_specialize=['T']) -+def causal_conv1d_states_fwd_kernel( -+ x, -+ initial_state, -+ final_state, -+ cu_seqlens, -+ T, -+ D, -+ W, -+ BD: tl.constexpr, -+ BW: tl.constexpr, -+ USE_INITIAL_STATE: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+): -+ i_d, i_n = tl.program_id(0), tl.program_id(1) -+ if IS_VARLEN: -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) -+ T = eos - bos -+ else: -+ bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) -+ -+ o_t = eos - BW + tl.arange(0, BW) -+ o_d = i_d * BD + tl.arange(0, BD) -+ o_w = W - BW + tl.arange(0, BW) -+ m_t = o_t >= tl.maximum(bos, eos - W) -+ m_d = o_d < D -+ m_w = (o_w >= 0) & (o_w < W) -+ -+ b_x = tl.load(x + o_t * D + o_d[:, None], mask=(m_t & m_d[:, None]), other=0) -+ if USE_INITIAL_STATE: -+ if T < BW: -+ o_c = W - (BW - T) + tl.arange(0, BW) -+ m_c = (o_c >= 0) & (o_c < W) -+ b_cache = tl.load(initial_state + i_n * D * W + o_d[:, None] * W + o_c, mask=m_d[:, None] & m_c, other=0) -+ b_x += b_cache -+ -+ tl.store(final_state + i_n * D * W + o_d[:, None] * W + o_w, b_x, mask=m_d[:, None] & m_w) -+ -+ -+@input_guard -+def causal_conv1d_update_states( -+ x: torch.Tensor, -+ state_len: int, -+ initial_state: Optional[torch.Tensor] = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+) -> torch.Tensor: -+ B, T, D, W = *x.shape, state_len -+ N = len(cu_seqlens) - 1 if cu_seqlens is not None else B -+ -+ final_state = torch.empty(N, D, W, dtype=x.dtype, device=x.device) -+ BD = min(triton.next_power_of_2(D), 256) -+ BW = W -+ grid = (triton.cdiv(D, BD), N) -+ causal_conv1d_states_fwd_kernel[grid]( -+ x=x, -+ initial_state=initial_state, -+ final_state=final_state, -+ cu_seqlens=cu_seqlens, -+ T=T, -+ D=D, -+ W=W, -+ BW=BW, -+ BD=BD, -+ ) -+ return final_state -+ -+ -+@triton.jit() -+def causal_conv1d_update_kernel_bdt_fwd( -+ x_ptr, # [B, D, T] -+ conv_state_ptr, # [B, D, ST] -+ conv_state_update_ptr, -+ weight_ptr, # [D, W] -+ bias_ptr, -+ conv_state_indices_ptr, -+ out_ptr, # [B, D, out_len] -+ batch: tl.constexpr, -+ dim: tl.constexpr, -+ state_len: tl.constexpr, # ST -+ seq_len: tl.constexpr, # T -+ width: tl.constexpr, # W -+ out_len: tl.constexpr, # output time -+ x_batch_stride: tl.constexpr, -+ conv_batch_stride: tl.constexpr, -+ out_batch_stride: tl.constexpr, -+ HAS_BIAS: tl.constexpr, -+ SILU_ACTIVATION: tl.constexpr, -+ T_CHK_SIZE: tl.constexpr, -+ D_CHK_SIZE: tl.constexpr, -+ NUM_T_CHK: tl.constexpr, -+ NUM_D_CHK: tl.constexpr, -+ ST_STORE_HEAD_TILE_SIZE: tl.constexpr, -+): -+ pid = tl.program_id(0) -+ pnum = tl.num_programs(0) -+ -+ total_task = batch * NUM_D_CHK * NUM_T_CHK -+ -+ for task_id in tl.range(pid, total_task, pnum): -+ di = task_id % NUM_D_CHK -+ bti = task_id // NUM_D_CHK -+ bi = bti // NUM_T_CHK -+ ti = bti % NUM_T_CHK -+ -+ w = tl.load( -+ tl.make_block_ptr( -+ weight_ptr, -+ shape=(dim, width), -+ strides=(width, 1), -+ offsets=(di * D_CHK_SIZE, 0), -+ block_shape=(D_CHK_SIZE, width), -+ order=(1, 0), -+ ), -+ boundary_check=(0, 1), -+ padding_option="zero", -+ ) -+ -+ if ti == 0: -+ st_b = tl.load( -+ tl.make_block_ptr( -+ conv_state_ptr + bi * state_len * dim, -+ shape=(dim, state_len), -+ strides=(state_len, 1), -+ offsets=(di * D_CHK_SIZE, state_len - (width - 1)), -+ block_shape=(D_CHK_SIZE, (width - 1) + T_CHK_SIZE), -+ order=(1, 0), -+ ), -+ boundary_check=(0, 1), -+ padding_option="zero", -+ ) -+ offset0_x = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE) -+ offset1_x = ti * T_CHK_SIZE + tl.arange(0, T_CHK_SIZE) -+ mask_x = (offset0_x < dim)[:, None] & ((offset1_x >= 0) & (offset1_x < seq_len))[None, :] -+ block_off_x = bi * dim * seq_len + offset0_x[:, None] * seq_len + offset1_x[None, :] -+ x_b_tmp = tl.load(x_ptr + block_off_x, mask=mask_x, other=0) -+ x_b = tl.insert_slice(st_b, x_b_tmp, (0, width - 1), (D_CHK_SIZE, T_CHK_SIZE), (1, 1)) -+ else: -+ offset0 = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE) -+ offset1 = ti * T_CHK_SIZE - (width - 1) + tl.arange(0, T_CHK_SIZE + width - 1) -+ mask = (offset0 < dim)[:, None] & ((offset1 >= 0) & (offset1 < seq_len))[None, :] -+ block_off = bi * dim * seq_len + offset0[:, None] * seq_len + offset1[None, :] -+ x_b = tl.load(x_ptr + block_off, mask=mask, other=0) -+ -+ out_block = tl.zeros((T_CHK_SIZE, D_CHK_SIZE), dtype=x_ptr.dtype.element_ty) -+ x_b = tl.trans(x_b, (1, 0)) -+ w = tl.trans(w, (1, 0)) -+ -+ new_state_start_off = seq_len - state_len -+ t_start_off = ti * T_CHK_SIZE - (width - 1) -+ t_end_off = (ti + 1) * T_CHK_SIZE -+ if t_end_off >= new_state_start_off: -+ t_off = t_start_off - new_state_start_off -+ if t_off < -(width - 1): -+ # NOTE: In order to avoid use tl.maximum for negative offset, -+ # we pre-compute a fix head tile size (ST_STORE_HEAD_TILE_SIZE) -+ # to store the scene of negative address -+ x_new_h = tl.extract_slice(x_b, (-t_off, 0), (ST_STORE_HEAD_TILE_SIZE, D_CHK_SIZE), (1, 1)) -+ x_new_h = tl.trans(x_new_h, (1, 0)) -+ nst_off_y0 = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE)[:, None] -+ nst_off_y1_h = tl.arange(0, ST_STORE_HEAD_TILE_SIZE)[None, :] -+ nst_mask_h = (nst_off_y0 < dim) & (nst_off_y1_h >= 0) & (nst_off_y1_h < state_len) -+ block_ptr_h = bi * dim * state_len + nst_off_y0 * state_len + nst_off_y1_h -+ tl.store(conv_state_update_ptr + block_ptr_h, x_new_h, mask=nst_mask_h) -+ else: -+ x_new_s = tl.extract_slice(x_b, (width - 1, 0), (T_CHK_SIZE, D_CHK_SIZE), (1, 1)) -+ x_new_s = tl.trans(x_new_s, (1, 0)) -+ nst_off_y0 = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE)[:, None] -+ nst_off_y1 = width - 1 + t_off + tl.arange(0, T_CHK_SIZE)[None, :] -+ nst_mask = (nst_off_y0 < dim) & (nst_off_y1 >= 0) & (nst_off_y1 < state_len) -+ block_ptr = bi * dim * state_len + nst_off_y0 * state_len + nst_off_y1 -+ tl.store(conv_state_update_ptr + block_ptr, x_new_s, mask=nst_mask) -+ -+ for owi in tl.range(0, width): -+ new_x = tl.extract_slice(x_b, (owi, 0), (T_CHK_SIZE, D_CHK_SIZE), (1, 1)) -+ w_chl_wi = tl.extract_slice(w, (owi, 0), (1, D_CHK_SIZE), (1, 1)) -+ x_mul_chl_wi = new_x * w_chl_wi -+ out_block += x_mul_chl_wi -+ out_block = tl.trans(out_block, (1, 0)) -+ -+ if SILU_ACTIVATION: -+ out_block = out_block * tl.sigmoid(out_block) -+ tl.store( -+ tl.make_block_ptr( -+ out_ptr, -+ shape=(batch, dim, out_len), -+ strides=(dim * out_len, out_len, 1), -+ offsets=(bi, di * D_CHK_SIZE, ti * T_CHK_SIZE), -+ block_shape=(1, D_CHK_SIZE, T_CHK_SIZE), -+ order=(2, 1, 0), -+ ), -+ out_block[None, :, :], -+ boundary_check=(0, 1, 2), -+ ) -+ -+ -+@input_guard -+def causal_conv1d_update_bdt_impl( -+ x: torch.Tensor, -+ conv_state: torch.Tensor, -+ weight: torch.Tensor, -+ bias: Optional[torch.Tensor] = None, -+ activation: Optional[str] = None, -+ conv_state_indices: Optional[str] = None, -+): -+ if isinstance(activation, bool): -+ activation = "silu" if activation is True else None -+ elif activation is not None: -+ if activation not in ["silu", "swish"]: -+ raise ValueError("activation must be one of 'silu' or 'swish'.") -+ unsqueeze = x.dim() == 2 -+ if unsqueeze: -+ x = x.unsqueeze(-1) -+ batch, dim, seqlen = x.shape -+ _, width = weight.shape -+ out = torch.empty_like(x) -+ -+ NUM_CORES = get_vector_num() -+ T_CHK_SIZE = 256 -+ D_CHK_SIZE = 16 -+ -+ if T_CHK_SIZE < width: -+ raise ValueError("T_CHK_SIZE must be >= width.") -+ -+ NUM_T_CHK = triton.cdiv(out.shape[-1], T_CHK_SIZE) -+ NUM_D_CHK = triton.cdiv(dim, D_CHK_SIZE) -+ conv_state_update = torch.empty_like(conv_state) -+ -+ # A const tile size variable to update negative address of conv state -+ ST_STORE_HEAD_TILE_SIZE = width if (seqlen % T_CHK_SIZE) > width else (width - seqlen % T_CHK_SIZE) % T_CHK_SIZE -+ causal_conv1d_update_kernel_bdt_fwd[(NUM_CORES, 1)]( -+ x, -+ conv_state, -+ conv_state_update, -+ weight, -+ bias, -+ conv_state_indices, -+ out, -+ batch=int(batch), -+ dim=int(dim), -+ state_len=int(conv_state.shape[-1]), -+ seq_len=int(x.shape[-1]), -+ width=int(width), -+ out_len=int(out.shape[-1]), -+ x_batch_stride=x.stride()[0], -+ conv_batch_stride=conv_state.stride()[0], -+ out_batch_stride=out.stride()[0], -+ HAS_BIAS=bias is not None, -+ SILU_ACTIVATION=activation in ["silu", "swish"], -+ T_CHK_SIZE=T_CHK_SIZE, -+ D_CHK_SIZE=D_CHK_SIZE, -+ NUM_T_CHK=NUM_T_CHK, -+ NUM_D_CHK=NUM_D_CHK, -+ ST_STORE_HEAD_TILE_SIZE=int(ST_STORE_HEAD_TILE_SIZE), -+ ) -+ conv_state.copy_(conv_state_update) -+ if unsqueeze: -+ out = out.squeeze(-1) -+ return out -diff --git a/megatron/core/ssm/triton/cumsum.py b/megatron/core/ssm/triton/cumsum.py -new file mode 100644 -index 000000000..805f56721 ---- /dev/null -+++ b/megatron/core/ssm/triton/cumsum.py -@@ -0,0 +1,144 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+from typing import Optional -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices -+ -+ -+@triton.heuristics({ -+ 'HAS_SCALE': lambda args: args['scale'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.jit(do_not_specialize=['T']) -+def chunk_local_cumsum_scalar_kernel( -+ s, -+ o, -+ scale, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ B: tl.constexpr, -+ H: tl.constexpr, -+ BLOCK_T: tl.constexpr, -+ REVERSE: tl.constexpr, -+ HAS_SCALE: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ HEAD_FIRST: tl.constexpr, -+ CHUNK_SIZE: tl.constexpr = 64, -+): -+ i_block, i_b = tl.program_id(0), tl.program_id(1) -+ N_CHUNKS: tl.constexpr = BLOCK_T // CHUNK_SIZE -+ -+ if IS_VARLEN: -+ i_s, i_block = tl.load(chunk_indices + i_block * 2).to(tl.int32), tl.load( -+ chunk_indices + i_block * 2 + 1 -+ ).to(tl.int32) -+ -+ bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load( -+ cu_seqlens + i_s + 1 -+ ).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ -+ ptr_s = tl.make_block_ptr( -+ s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) -+ ) -+ ptr_o = tl.make_block_ptr( -+ o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) -+ ) -+ b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32) -+ b_s = tl.reshape(b_s, (N_CHUNKS, CHUNK_SIZE, H)) -+ b_s = tl.trans(b_s, (1, 0, 2)) -+ b_o = tl.cumsum(b_s, axis=0) -+ if REVERSE: -+ b_z = tl.sum(b_s, axis=0) -+ b_o = -b_o + b_z[None] + b_s -+ if HAS_SCALE: -+ b_o *= scale -+ b_o = tl.trans(b_o, (1, 0, 2)) -+ b_o = tl.reshape(b_o, (BLOCK_T, H)) -+ -+ tl.store(ptr_o, b_o.to(ptr_o.dtype.element_ty), boundary_check=(0,)) -+ return -+ -+ -+def chunk_local_cumsum_scalar( -+ g: torch.Tensor, -+ chunk_size: int, -+ reverse: bool = False, -+ scale: float = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ head_first: bool = False, -+ output_dtype: Optional[torch.dtype] = torch.float -+) -> torch.Tensor: -+ -+ B, T, H = g.shape -+ if chunk_size != 2 ** (chunk_size.bit_length() - 1): -+ raise ValueError( -+ f"chunk_size must be a power of 2, chunk_size is{chunk_size}" -+ ) -+ # We adjust the tiling strategy to prevent overflow in in backward passes and context parallel scenarios -+ # while maximizing UB utilization where possible. -+ # The tiling strategy is as follows: -+ # 1. BT must be greater than or equal to chunk_size. -+ # 2. UB estimation varies directly with H. -+ # 3. BT in reverse mode is smaller than in forward mode. -+ BT = max(chunk_size, triton.next_power_of_2((1 << 11 if reverse else 1 << 12) // H)) -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) -+ grid = (NT, B) -+ chunk_local_cumsum_scalar_kernel[grid]( -+ s=g_org, -+ o=g, -+ scale=scale, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ B=B, -+ H=H, -+ BLOCK_T=BT, -+ HEAD_FIRST=head_first, -+ REVERSE=reverse, -+ CHUNK_SIZE=chunk_size, -+ ) -+ return g -+ -+ -+def chunk_local_cumsum( -+ g: torch.Tensor, -+ chunk_size: int, -+ reverse: bool = False, -+ scale: float = None, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ head_first: bool = False, -+ output_dtype: Optional[torch.dtype] = torch.float, -+ **kwargs -+) -> torch.Tensor: -+ if cu_seqlens is not None: -+ if g.shape[0] != 1: -+ raise ValueError( -+ f"Only batch size 1 is supported when cu_seqlens are provided, current size is{g.shape[0]}" -+ ) -+ if len(g.shape) == 3: -+ return chunk_local_cumsum_scalar( -+ g=g, -+ chunk_size=chunk_size, -+ reverse=reverse, -+ scale=scale, -+ cu_seqlens=cu_seqlens, -+ head_first=head_first, -+ output_dtype=output_dtype -+ ) -+ else: -+ raise ValueError( -+ f"Unsupported input shape {g.shape}, " -+ f"which should be (B, T, H, D) if `head_first=False` " -+ f"or (B, H, T, D) otherwise" -+ ) -diff --git a/megatron/core/ssm/triton/l2norm.py b/megatron/core/ssm/triton/l2norm.py -new file mode 100644 -index 000000000..64b85c241 ---- /dev/null -+++ b/megatron/core/ssm/triton/l2norm.py -@@ -0,0 +1,315 @@ -+# -*- coding: utf-8 -*- -+# Copyright c) 2023-2025 Songlin Yang Yu Zhang -+# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved. -+ -+from typing import Optional -+ -+import torch -+import torch.nn as nn -+import triton -+import triton.language as tl -+ -+from .utils import input_guard, is_amd -+ -+BT_LIST = [8, 16, 32, 64, 128] -+NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if is_amd else [1, 2, 4, 8, 16, 32] -+ -+ -+@triton.autotune( -+ configs=[ -+ triton.Config({}, num_warps=num_warps) -+ for num_warps in NUM_WARPS_AUTOTUNE -+ ], -+ key=['D'] -+) -+@triton.jit -+def l2norm_fwd_kernel1( -+ x, -+ y, -+ rstd, -+ eps, -+ D, -+ BD: tl.constexpr, -+): -+ i_t = tl.program_id(0) -+ x += i_t * D -+ y += i_t * D -+ # Compute mean and variance -+ cols = tl.arange(0, BD) -+ mask = cols < D -+ -+ b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) -+ b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) -+ b_y = b_x * b_rstd -+ tl.store(y + cols, b_y, mask=mask) -+ tl.store(rstd + i_t, b_rstd) -+ -+ -+@triton.autotune( -+ configs=[ -+ triton.Config({}, num_warps=num_warps) -+ for num_warps in NUM_WARPS_AUTOTUNE -+ ], -+ key=['D'] -+) -+@triton.jit -+def l2norm_bwd_kernel1( -+ y, -+ rstd, -+ dy, -+ dx, -+ eps, -+ D, -+ BD: tl.constexpr, -+): -+ i_t = tl.program_id(0) -+ y += i_t * D -+ dx += i_t * D -+ dy += i_t * D -+ -+ cols = tl.arange(0, BD) -+ mask = cols < D -+ b_y = tl.load(y + cols, mask=mask, other=0.0).to(tl.float32) -+ b_rstd = tl.load(rstd + i_t).to(tl.float32) -+ b_dy = tl.load(dy + cols, mask=mask, other=0.0).to(tl.float32) -+ b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd -+ tl.store(dx + cols, b_dx, mask=mask) -+ -+ -+@triton.autotune( -+ configs=[ -+ triton.Config({'BT': BT}, num_warps=num_warps) -+ for num_warps in [1, 2, 4, 8, 16] -+ for BT in BT_LIST -+ ], -+ key=['D', 'NB'] -+) -+@triton.jit -+def l2norm_fwd_kernel( -+ x, -+ y, -+ rstd, -+ eps, -+ T: tl.constexpr, -+ D: tl.constexpr, -+ BD: tl.constexpr, -+ NB: tl.constexpr, -+ BT: tl.constexpr, -+ bt_size, -+): -+ i_t = tl.program_id(0) -+ for offset in range(0, bt_size): -+ block_start = (i_t * bt_size + offset) * BT -+ if block_start < T: -+ p_x = tl.make_block_ptr(x, (T, D), (D, 1), (block_start, 0), (BT, BD), (1, 0)) -+ p_y = tl.make_block_ptr(y, (T, D), (D, 1), (block_start, 0), (BT, BD), (1, 0)) -+ p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (block_start,), (BT,), (0,)) -+ -+ b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) -+ b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) -+ b_y = b_x * b_rstd[:, None] -+ -+ tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) -+ tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) -+ -+ -+@triton.autotune( -+ configs=[ -+ triton.Config({'BT': BT}, num_warps=num_warps) -+ for num_warps in [1, 2, 4, 8, 16] -+ for BT in BT_LIST -+ ], -+ key=['D', 'NB'] -+) -+@triton.jit -+def l2norm_bwd_kernel( -+ y, -+ rstd, -+ dy, -+ dx, -+ eps, -+ T: tl.constexpr, -+ D: tl.constexpr, -+ BD: tl.constexpr, -+ NB: tl.constexpr, -+ BT: tl.constexpr, -+ bt_size, -+): -+ i_t_start = tl.program_id(0) -+ num_blocks = bt_size -+ -+ total_i_t = tl.cdiv(T, BT) -+ base_tasks_per_block = total_i_t // num_blocks -+ remainder_tasks = total_i_t % num_blocks -+ -+ if i_t_start < remainder_tasks: -+ tasks_this_block = base_tasks_per_block + 1 -+ start_i_t = i_t_start * tasks_this_block -+ else: -+ tasks_this_block = base_tasks_per_block -+ start_i_t = i_t_start * base_tasks_per_block + remainder_tasks -+ -+ for task_idx in range(tasks_this_block): -+ i_t = start_i_t + task_idx -+ block_start = i_t * BT -+ if block_start < T: -+ p_y = tl.make_block_ptr(y, (T, D), (D, 1), (block_start, 0), (BT, BD), (1, 0)) -+ p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (block_start,), (BT,), (0,)) -+ p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (block_start, 0), (BT, BD), (1, 0)) -+ p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (block_start, 0), (BT, BD), (1, 0)) -+ -+ b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) -+ b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) -+ b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) -+ b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] -+ tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+def l2norm_fwd( -+ x: torch.Tensor, -+ eps: float = 1e-6, -+ output_dtype: Optional[torch.dtype] = None -+): -+ x_shape_og = x.shape -+ x = x.view(-1, x.shape[-1]) -+ # allocate output -+ if output_dtype is None: -+ y = torch.empty_like(x) -+ else: -+ y = torch.empty_like(x, dtype=output_dtype) -+ assert y.stride(-1) == 1 -+ T, D = x.shape[0], x.shape[-1] -+ # Less than 64KB per feature: enqueue fused kernel -+ MAX_FUSED_SIZE = 65536 // x.element_size() -+ BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) -+ if D > BD: -+ raise RuntimeError("This layer doesn't support feature dim >= 64KB.") -+ -+ rstd = torch.empty((T,), dtype=torch.float32, device=x.device) -+ if D <= 512: -+ NB = triton.cdiv(T, 2048) -+ bt_size = 32 -+ -+ def grid(meta): -+ new_bt = meta['BT'] * bt_size -+ return (triton.cdiv(T, new_bt), ) -+ -+ l2norm_fwd_kernel[grid]( -+ x=x, -+ y=y, -+ rstd=rstd, -+ eps=eps, -+ T=T, -+ D=D, -+ BD=BD, -+ NB=NB, -+ bt_size=bt_size, -+ ) -+ else: -+ l2norm_fwd_kernel1[(T,)]( -+ x=x, -+ y=y, -+ rstd=rstd, -+ eps=eps, -+ D=D, -+ BD=BD, -+ ) -+ return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) -+ -+ -+def l2norm_bwd( -+ y: torch.Tensor, -+ rstd: torch.Tensor, -+ dy: torch.Tensor, -+ eps: float = 1e-6 -+): -+ y_shape_og = y.shape -+ y = y.view(-1, dy.shape[-1]) -+ dy = dy.view(-1, dy.shape[-1]) -+ assert dy.shape == y.shape -+ # allocate output -+ dx = torch.empty_like(y) -+ T, D = y.shape[0], y.shape[-1] -+ # Less than 64KB per feature: enqueue fused kernel -+ MAX_FUSED_SIZE = 65536 // y.element_size() -+ BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) -+ if D > BD: -+ raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") -+ -+ if D <= 512: -+ NB = triton.cdiv(T, 2048) -+ bt_size = 40 -+ l2norm_bwd_kernel[(bt_size,)]( -+ y=y, -+ rstd=rstd, -+ dy=dy, -+ dx=dx, -+ eps=eps, -+ T=T, -+ D=D, -+ BD=BD, -+ NB=NB, -+ bt_size=bt_size, -+ ) -+ else: -+ l2norm_bwd_kernel1[(T,)]( -+ y=y, -+ rstd=rstd, -+ dy=dy, -+ dx=dx, -+ eps=eps, -+ D=D, -+ BD=BD, -+ ) -+ -+ return dx.view(y_shape_og) -+ -+ -+class L2NormFunction(torch.autograd.Function): -+ -+ @staticmethod -+ @input_guard -+ def forward( -+ ctx, -+ x, -+ eps=1e-6, -+ output_dtype=None -+ ): -+ y, rstd = l2norm_fwd(x, eps, output_dtype) -+ ctx.eps = eps -+ ctx.x_dtype = x.dtype -+ ctx.save_for_backward(y, rstd) -+ return y -+ -+ @staticmethod -+ @input_guard -+ def backward(ctx, dy): -+ y, rstd = ctx.saved_tensors -+ dx = l2norm_bwd(y, rstd, dy, ctx.eps) -+ return dx, None, None -+ -+ -+def l2norm( -+ x: torch.Tensor, -+ eps: float = 1e-6, -+ output_dtype: Optional[torch.dtype] = None -+) -> torch.Tensor: -+ return L2NormFunction.apply(x, eps, output_dtype) -+ -+ -+l2_norm = l2norm -+ -+ -+class L2Norm(nn.Module): -+ -+ def __init__( -+ self, -+ eps: float = 1e-6, -+ output_dtype: Optional[torch.dtype] = None -+ ): -+ super().__init__() -+ self.eps = eps -+ self.output_dtype = output_dtype -+ -+ def forward(self, x: torch.Tensor) -> torch.Tensor: -+ return l2norm(x, self.eps, self.output_dtype) -diff --git a/megatron/core/ssm/triton/solve_tril.py b/megatron/core/ssm/triton/solve_tril.py -new file mode 100644 -index 000000000..be56bfb59 ---- /dev/null -+++ b/megatron/core/ssm/triton/solve_tril.py -@@ -0,0 +1,519 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+# Copyright (c) 2023-2025, By Triton_Ascend & sglang_ascend -+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. -+ -+import os -+from typing import Optional -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices, make_tensor_descriptor, input_guard, is_amd -+ -+ -+def _ensure_slice_ops() -> bool: -+ """Probe and attach tl.extract_slice / insert_slice if missing; return success.""" -+ if hasattr(tl, "extract_slice") and hasattr(tl, "insert_slice"): -+ return True -+ try: -+ from triton.language.extra.cann.extension import extract_slice, insert_slice -+ tl.extract_slice = extract_slice -+ tl.insert_slice = insert_slice -+ return True -+ except ImportError: -+ return False -+ -+_TRITON_SLICE_AVAILABLE: bool = _ensure_slice_ops() -+FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') -+ -+ -+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -+@triton.jit(do_not_specialize=["T"]) -+def solve_tril_16x16_loop_kernel_paral_v3( -+ A_ptr, -+ Ad_ptr, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ H: tl.constexpr, -+ BT: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ LARGE_BLOCK_T: tl.constexpr, -+ NT: tl.constexpr, -+ BH: tl.constexpr, -+): -+ worker_id = tl.program_id(0) -+ total_tasks = NT * BH -+ num_tasks = total_tasks // 48 -+ remainder = total_tasks - num_tasks * 48 -+ upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder)) -+ lower_bound = num_tasks * worker_id + min(worker_id, remainder) -+ for task_id in range(lower_bound, upper_bound): -+ i_t = task_id // BH -+ i_bh = task_id % BH -+ i_b, i_h = i_bh // H, i_bh % H -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( -+ chunk_indices + i_t * 2 + 1 -+ ).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( -+ cu_seqlens + i_n + 1 -+ ).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ -+ A = A_ptr + (bos * H + i_h) * BT -+ Ad = Ad_ptr + (bos * H + i_h) * 16 -+ -+ base_t = i_t * LARGE_BLOCK_T -+ -+ NTASKS: tl.constexpr = 2 -+ N_BLOCKS: tl.constexpr = LARGE_BLOCK_T // 16 // NTASKS -+ -+ for taskid in range(0, NTASKS): -+ base_t += taskid * (LARGE_BLOCK_T // NTASKS) -+ -+ b_A = tl.zeros((N_BLOCKS, 16, 16), dtype=tl.float32) # (N_BLOCKS, 16, 16) -+ for blkid in range(0, N_BLOCKS): -+ row_start_o = base_t + blkid * 16 -+ col_start_o = row_start_o % BT -+ # using ptr with mask instead of tl.load(block_ptr) -+ offs_rows_in_block = tl.arange(0, 16) -+ offs_cols_in_block = tl.arange(0, 16) -+ ptr_A_subrec16 = ( -+ A -+ + row_start_o * H * BT -+ + col_start_o -+ + offs_rows_in_block[:, None] * H * BT -+ + offs_cols_in_block[None, :] -+ ) -+ global_rows = row_start_o + offs_rows_in_block[:, None] -+ global_cols = col_start_o + offs_cols_in_block[None, :] -+ load_mask = (global_rows < T) & (global_cols < BT) -+ b_A_subrec16 = tl.load(ptr_A_subrec16, mask=load_mask, other=0.0).to( -+ tl.float32 -+ ) -+ b_A = tl.insert_slice( -+ ful=b_A, -+ sub=b_A_subrec16[None, :, :], # (1, 16, 16) -+ offsets=[blkid, 0, 0], -+ sizes=[1, 16, 16], -+ strides=[1, 1, 1], -+ ) -+ -+ # load multi 16x16 -+ local_ori_A = tl.trans(b_A, (1, 0, 2)) -+ local_ori_A = tl.reshape(local_ori_A, (16, 16 * N_BLOCKS)) # (16, N_BLOCKS*16) -+ -+ # change mask into matrix elementwise action -+ tmp = tl.arange(0, 16).to(tl.float32) -+ rows = tmp[:, None] -+ cols = tmp[None, :] -+ is_lower = (rows > cols).to(b_A.dtype) -+ b_A = -b_A * is_lower -+ -+ for i in range(1, 16): -+ nblks_vec16 = -tl.extract_slice( -+ local_ori_A, (i, 0), (1, 16 * N_BLOCKS), (16 * N_BLOCKS, 1) -+ ) -+ b_a = tl.reshape(nblks_vec16, (N_BLOCKS, 16)) -+ -+ dot_tmp = tl.trans(b_a[:, :, None] * b_A, (1, 0, 2)) -+ dot_product = tl.sum(dot_tmp, 0) -+ b_a = b_a + dot_product # (N_BLOCKS, 16) -+ -+ b_a_new_expanded = b_a[:, None, :] # (N_BLOCKS, 1, 16) -+ b_A = tl.insert_slice( -+ ful=b_A, -+ sub=b_a_new_expanded, -+ offsets=[0, i, 0], -+ sizes=[N_BLOCKS, 1, 16], -+ strides=[1, 1, 1], -+ ) -+ -+ on_diagonal = rows == cols -+ b_A = tl.where(on_diagonal, b_A + 1.0, b_A) -+ -+ b_A = tl.reshape(b_A, (N_BLOCKS * 16, 16)) -+ # using ptr with mask instead of tl.load(block_ptr) -+ offs_rows_to_store = tl.arange(0, N_BLOCKS * 16) -+ offs_cols_to_store = tl.arange(0, 16) -+ p_Ai = ( -+ Ad -+ + base_t * H * 16 -+ + 0 -+ + offs_rows_to_store[:, None] * H * 16 -+ + offs_cols_to_store[None, :] -+ ) -+ global_store_rows = base_t + offs_rows_to_store[:, None] -+ store_mask = global_store_rows < T -+ tl.store( -+ p_Ai, -+ b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), -+ mask=store_mask, -+ ) -+ -+ -+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -+@triton.jit(do_not_specialize=["T", "NT"]) -+def merge_16x16_to_32x32_loop_inverse_kernel( -+ A, -+ Ad, -+ Ai, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ NT, -+ H: tl.constexpr, -+ BT: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ BH: tl.constexpr, -+): -+ worker_id = tl.program_id(0) -+ total_tasks = NT * BH -+ num_tasks = total_tasks // 24 -+ remainder = total_tasks - num_tasks * 24 -+ upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder)) -+ lower_bound = num_tasks * worker_id + min(worker_id, remainder) -+ for task_id in range(lower_bound, upper_bound): -+ i_tt = task_id // BH -+ i_bh = task_id % BH -+ i_b, i_h = i_bh // H, i_bh % H -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_tt * 2).to(tl.int32), tl.load( -+ chunk_indices + i_tt * 2 + 1 -+ ).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( -+ cu_seqlens + i_n + 1 -+ ).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ i_t = i_tt -+ -+ A_ptr = A + (bos * H + i_h) * BT -+ Ad_ptr = Ad + (bos * H + i_h) * 16 -+ Ai_ptr = Ai + (bos * H + i_h) * 32 -+ -+ p_A_21 = tl.make_block_ptr( -+ A_ptr, (T, BT), (H * BT, 1), (i_t * 32 + 16, 0 + i_t % (BT // 32) * 32), (16, 16), (1, 0) -+ ) -+ p_Ad_11 = tl.make_block_ptr( -+ Ad_ptr, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0) -+ ) -+ p_Ad_22 = tl.make_block_ptr( -+ Ad_ptr, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0) -+ ) -+ p_Ai_11 = tl.make_block_ptr( -+ Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0) -+ ) -+ p_Ai_22 = tl.make_block_ptr( -+ Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0) -+ ) -+ p_Ai_21 = tl.make_block_ptr( -+ Ai_ptr, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0) -+ ) -+ -+ A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) -+ Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) -+ Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) -+ Ai_21 = -tl.dot( -+ tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee" -+ ) -+ tl.store( -+ p_Ai_11, -+ Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ tl.store( -+ p_Ai_22, -+ Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ tl.store( -+ p_Ai_21, -+ Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ -+ -+@triton.heuristics( -+ { -+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, -+ } -+) -+@triton.jit(do_not_specialize=["T", "NT"]) -+def merge_32x32_to_64x64_loop_inverse_kernel( -+ A, -+ Ad, -+ Ai, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ NT, -+ H: tl.constexpr, -+ BT: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ BH: tl.constexpr, -+): -+ worker_id = tl.program_id(0) -+ total_tasks = NT * BH -+ num_tasks = total_tasks // 24 -+ remainder = total_tasks - num_tasks * 24 -+ upper_bound = min(total_tasks, num_tasks * (worker_id + 1) + min(worker_id + 1, remainder)) -+ lower_bound = num_tasks * worker_id + min(worker_id, remainder) -+ for task_id in range(lower_bound, upper_bound): -+ i_tt = task_id // BH -+ i_bh = task_id % BH -+ i_b, i_h = i_bh // H, i_bh % H -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_tt * 2).to(tl.int32), tl.load( -+ chunk_indices + i_tt * 2 + 1 -+ ).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( -+ cu_seqlens + i_n + 1 -+ ).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ i_t = i_tt -+ -+ A_ptr = A + (bos * H + i_h) * BT -+ Ad_ptr = Ad + (bos * H + i_h) * 32 -+ Ai_ptr = Ai + (bos * H + i_h) * 64 -+ -+ p_A_21 = tl.make_block_ptr( -+ A_ptr, (T, BT), (H * BT, 1), (i_t * 64 + 32, 0 + i_t % (BT // 64) * 64), (32, 32), (1, 0) -+ ) -+ -+ p_Ad_11 = tl.make_block_ptr( -+ Ad_ptr, (T, 32), (H * 32, 1), (i_t * 64, 0), (32, 32), (1, 0) -+ ) -+ p_Ad_22 = tl.make_block_ptr( -+ Ad_ptr, (T, 32), (H * 32, 1), (i_t * 64 + 32, 0), (32, 32), (1, 0) -+ ) -+ -+ p_Ai_11 = tl.make_block_ptr( -+ Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64, 0), (32, 32), (1, 0) -+ ) -+ p_Ai_22 = tl.make_block_ptr( -+ Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32), (32, 32), (1, 0) -+ ) -+ p_Ai_21 = tl.make_block_ptr( -+ Ai_ptr, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (32, 32), (1, 0) -+ ) -+ -+ A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) -+ Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32) -+ Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32) -+ Ai_21 = -tl.dot( -+ tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee" -+ ) -+ tl.store( -+ p_Ai_11, -+ Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ tl.store( -+ p_Ai_22, -+ Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ tl.store( -+ p_Ai_21, -+ Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), -+ boundary_check=(0, 1), -+ ) -+ -+ -+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -+@triton.jit(do_not_specialize=['T']) -+def solve_tril_64x64_kernel( -+ A, -+ Ai, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ H: tl.constexpr, -+ BT: tl.constexpr, -+ USE_TMA: tl.constexpr, -+ IS_VARLEN: tl.constexpr, -+ DOT_PRECISION: tl.constexpr -+): -+ i_t, i_bh = tl.program_id(0), tl.program_id(1) -+ i_b, i_h = i_bh // H, i_bh % H -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T = eos - bos -+ else: -+ bos, eos = i_b * T, i_b * T + T -+ o_i = tl.arange(0, 64) -+ m_I = o_i[:, None] == o_i[None, :] -+ -+ A = A + (bos * H + i_h) * BT -+ Ai = Ai + (bos * H + i_h) * 64 -+ -+ offset = (i_t * 64) % BT -+ if not USE_TMA: -+ p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * 64, offset), (64, 64), (1, 0)) -+ b_A = -tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) -+ else: -+ desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [64, 64]) -+ desc_o = make_tensor_descriptor(Ai, [T, 64], [H * 64, 1], [64, 64]) -+ b_A = -desc.load([i_t * 64, offset]).to(tl.float32) -+ -+ for i in range(2, min(64, T - i_t * 64)): -+ b_a = -tl.load(A + (i_t * 64 + i) * H * BT + o_i + offset) -+ b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) -+ b_A = tl.where((o_i == i)[:, None], b_a, b_A) -+ b_A += m_I -+ if not USE_TMA: -+ p_Ai = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (64, 64), (1, 0)) -+ tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) -+ else: -+ desc_o.store([i_t * 64, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) -+ -+ -+def solve_tril_64( -+ A: torch.Tensor, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ output_dtype: torch.dtype = torch.float, -+ ): -+ B, T, H, BT = A.shape -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) -+ -+ Ai = torch.zeros_like(A, dtype=output_dtype) -+ solve_tril_64x64_kernel[NT, B * H]( -+ A=A, -+ Ai=Ai, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ BT=BT, -+ USE_TMA=False, -+ DOT_PRECISION=FLA_TRIL_PRECISION, -+ ) -+ return Ai -+ -+ -+@input_guard -+def solve_tril( -+ A: torch.Tensor, -+ cu_seqlens: Optional[torch.Tensor] = None, -+ output_dtype: torch.dtype = torch.float -+) -> torch.Tensor: -+ """ -+ Compute the inverse of the matrix I + A -+ A should be strictly lower triangular, i.e., A.triu() == 0. -+ -+ Args: -+ A (torch.Tensor): -+ [B, T, H, BT], where BT should only be 16, 32, or 64. -+ cu_seqlens (torch.Tensor): -+ The cumulative sequence lengths of the input tensor. Default: `None`. -+ output_dtype (torch.dtype): -+ The dtype of the output tensor. Default: `torch.float`. -+ If `None`, the output dtype will be the same as the input dtype. -+ -+ Returns: -+ (I + A)^-1 with the same shape as A -+ """ -+ output_dtype = A.dtype if output_dtype is None else output_dtype -+ if not _TRITON_SLICE_AVAILABLE: -+ if A.shape[-1] not in [64]: -+ raise ValueError( -+ f"A shape BT should in [64], but current is {A.shape[-1]}" -+ ) -+ return solve_tril_64(A, cu_seqlens, output_dtype) -+ if A.shape[-1] not in [16, 32, 64]: -+ raise ValueError( -+ f"A shape BT should in [16, 32, 64], but current is {A.shape[-1]}" -+ ) -+ -+ B, T, H, BT = A.shape -+ # If BT matches the current processing level (final step), use output_dtype -+ # (e.g. BF16) so the kernel can downcast internally, avoiding an extra -+ # external cast that hurts performance. Otherwise, keep FP32 to preserve -+ # precision for subsequent computation stages. -+ Ad = torch.empty( -+ B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype -+ ) -+ -+ LARGE_BLOCK_T = 608 * 2 -+ -+ chunk_indices = ( -+ prepare_chunk_indices(cu_seqlens, LARGE_BLOCK_T) -+ if cu_seqlens is not None -+ else None -+ ) -+ NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, LARGE_BLOCK_T) -+ solve_tril_16x16_loop_kernel_paral_v3[(48,)]( -+ A, -+ Ad, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ BT=BT, -+ LARGE_BLOCK_T=LARGE_BLOCK_T, -+ NT=NT, -+ BH=B * H, -+ ) -+ -+ if BT == 16: -+ return Ad -+ -+ # Same dtype logic as above: output_dtype for the final step, FP32 otherwise. -+ Ai = torch.zeros( -+ B, T, H, 32, device=A.device, dtype=torch.float if BT != 32 else output_dtype -+ ) -+ -+ chunk_indices = ( -+ prepare_chunk_indices(cu_seqlens, 32) if cu_seqlens is not None else None -+ ) -+ NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 32) -+ merge_16x16_to_32x32_loop_inverse_kernel[(24,)]( -+ A=A, -+ Ad=Ad, -+ Ai=Ai, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ BT=BT, -+ NT=NT, -+ BH=B * H, -+ ) -+ if BT == 32: -+ return Ai -+ -+ Ad = Ai -+ # Same dtype logic as above: output_dtype for the final step, FP32 otherwise. -+ Ai = torch.zeros( -+ B, T, H, 64, device=A.device, dtype=torch.float if BT != 64 else output_dtype -+ ) -+ chunk_indices = ( -+ prepare_chunk_indices(cu_seqlens, 64) if cu_seqlens is not None else None -+ ) -+ NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 64) -+ merge_32x32_to_64x64_loop_inverse_kernel[(24,)]( -+ A=A, -+ Ad=Ad, -+ Ai=Ai, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ H=H, -+ BT=BT, -+ NT=NT, -+ BH=B * H, -+ ) -+ if BT == 64: -+ return Ai -+ return Ai -diff --git a/megatron/core/ssm/triton/utils.py b/megatron/core/ssm/triton/utils.py -new file mode 100644 -index 000000000..fc5732596 ---- /dev/null -+++ b/megatron/core/ssm/triton/utils.py -@@ -0,0 +1,370 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+ -+# pylint: disable=no-name-in-module,consider-using-from-import,pointless-string-statement,redefined-outer-name -+ -+import itertools -+import contextlib -+import os -+import functools -+import warnings -+import logging -+from enum import Enum -+from functools import lru_cache -+from typing import Any, Callable, Optional -+from packaging import version -+ -+import torch -+import triton -+import triton.language as tl -+import triton.language.extra.libdevice as tldevice -+import triton.runtime.driver as driver -+ -+logger = logging.getLogger(__name__) -+ -+FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" -+ -+ -+def tensor_cache(fn: Optional[Callable[..., torch.Tensor]] = None, *, maxsize: int = 1) -> Any: -+ """ -+ A decorator that caches the most recent results of a function with tensor inputs. -+ -+ This decorator will store the outputs of the decorated function for the most recent -+ set of input tensors, up to `maxsize` entries. If the function is called again with -+ the same input tensors, it will return the cached result. -+ -+ When maxsize=1 (default), the behavior is identical to caching only the most recent result. -+ Can be used as @tensor_cache or @tensor_cache(maxsize=n). -+ -+ Args: -+ fn (Callable[..., torch.Tensor], optional): -+ The function to be decorated when used without parentheses. -+ maxsize (int): -+ Maximum number of input combinations to cache. Default is 1. -+ -+ Returns: -+ Callable[..., torch.Tensor]: -+ A wrapped version of the input function with caching. -+ """ -+ if maxsize < 1: -+ raise ValueError("maxsize must be at least 1") -+ -+ def _is_match(a: Any, b: Any) -> bool: -+ if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): -+ return a is b -+ try: -+ return a == b -+ except Exception: -+ return a is b -+ -+ def _make_wrapper(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: -+ cache: list = [] -+ -+ @functools.wraps(fn) -+ def wrapper(*args: Any, **kwargs: Any) -> Any: -+ for i, (cached_args, cached_kwargs, cached_result) in enumerate(cache): -+ if len(args) == len(cached_args) and len(kwargs) == len(cached_kwargs): -+ if all(_is_match(a, b) for a, b in zip(args, cached_args)) and all( -+ k in cached_kwargs and _is_match(v, cached_kwargs[k]) for k, v in kwargs.items() -+ ): -+ if i != 0: -+ cache.insert(0, cache.pop(i)) -+ return cached_result -+ -+ result = fn(*args, **kwargs) -+ cache.insert(0, (args, kwargs, result)) -+ if len(cache) > maxsize: -+ cache.pop() -+ return result -+ -+ return wrapper -+ -+ if fn is not None: -+ return _make_wrapper(fn) -+ return _make_wrapper -+ -+ -+@tensor_cache -+def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: -+ return cu_seqlens[1:] - cu_seqlens[:-1] -+ -+ -+@tensor_cache(maxsize=3) -+def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: -+ indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) -+ return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) -+ -+ -+def get_abs_err(x, y): -+ return (x.detach() - y.detach()).flatten().abs().max().item() -+ -+ -+def get_err_ratio(x, y): -+ err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item() -+ base = (x.detach()).flatten().square().mean().sqrt().item() -+ return err / (base + 1e-8) -+ -+ -+def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): -+ abs_atol = get_abs_err(ref, tri) -+ msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}" -+ logger.info(msg) -+ error_rate = get_err_ratio(ref, tri) -+ if abs_atol <= err_atol: -+ return -+ if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): -+ if error_rate > ratio: -+ warnings.warn(msg) -+ else: -+ assert error_rate < ratio, msg -+ -+ -+if hasattr(triton.language, '_experimental_make_tensor_descriptor'): -+ # For Triton 3.3.x -+ make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor -+elif hasattr(triton.language, 'make_tensor_descriptor'): -+ # For Triton 3.4.x and later -+ make_tensor_descriptor = triton.language.make_tensor_descriptor -+else: -+ """ -+ Fallback implementation when TMA is not supported. -+ Returns None to indicate TMA descriptors are unavailable. -+ Just make triton compiler happy. -+ """ -+ -+ @triton.jit -+ def make_tensor_descriptor( -+ base, -+ shape, -+ strides, -+ block_shape, -+ _builder=None, -+ ): -+ return None -+ -+ -+def _cpu_device_warning(): -+ warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=1) -+ -+ -+@lru_cache(maxsize=None) -+def get_available_device() -> str: -+ try: -+ return triton.runtime.driver.active.get_current_target().backend -+ except BaseException: -+ _cpu_device_warning() -+ return 'cpu' -+ -+ -+def map_triton_backend_to_torch_device() -> str: -+ backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ... -+ return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend) -+ -+ -+device = get_available_device() if get_available_device() != 'hip' else 'cuda' -+device_torch_lib = getattr(torch, device) -+device_platform = get_available_device() -+is_amd = device_platform == 'hip' -+is_nvidia = device_platform == 'cuda' -+is_nvidia_hopper = is_nvidia and ( -+ 'NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9 -+) -+ -+is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8 -+is_tma_supported = ( -+ (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) -+ and os.environ.get('FLA_NO_USE_TMA', '0') != '1' -+ and ( -+ hasattr(triton.language, '_experimental_make_tensor_descriptor') -+ or hasattr(triton.language, 'make_tensor_descriptor') -+ ) -+) -+ -+if is_nvidia and not is_tf32_supported: -+ # Make old card happy, since triton will use tf32 by default. -+ # This is a workaround for old nvidia card. -+ os.environ['TRITON_F32_DEFAULT'] = 'ieee' -+ -+ -+@lru_cache(maxsize=None) -+def check_pytorch_version(version_s: str = '2.4') -> bool: -+ return version.parse(torch.__version__) >= version.parse(version_s) -+ -+ -+if check_pytorch_version('2.4'): -+ device = 'cuda' if device == 'cpu' else device -+ autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) -+ autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) -+ -+ def custom_device_ctx(index: int): -+ return device_torch_lib.device(index) -+else: -+ assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.' -+ autocast_custom_fwd = device_torch_lib.amp.custom_fwd -+ autocast_custom_bwd = device_torch_lib.amp.custom_bwd -+ -+ def custom_device_ctx(index: int): -+ return torch.cuda.device(index) -+ -+ -+def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: -+ """ -+ A decorator to make sure all input tensors are contiguous and set the device based on input tensors. -+ """ -+ -+ @functools.wraps(fn) -+ def wrapper(*args, **kwargs): -+ contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args) -+ contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()} -+ -+ tensor = None -+ for arg in args: -+ if isinstance(arg, torch.Tensor): -+ tensor = arg -+ break -+ if tensor is None: -+ for value in kwargs.values(): -+ if isinstance(value, torch.Tensor): -+ tensor = value -+ break -+ -+ if tensor is not None: -+ ctx = custom_device_ctx(tensor.device.index) -+ else: -+ ctx = contextlib.nullcontext() -+ -+ with ctx: -+ return fn(*contiguous_args, **contiguous_kwargs) -+ -+ return wrapper -+ -+ -+@tensor_cache -+def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: -+ return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) -+ -+ -+if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': -+ exp = tldevice.fast_expf -+ exp2 = tldevice.exp2 -+ log = tldevice.fast_logf -+ log2 = tldevice.fast_log2f -+else: -+ exp = tl.exp -+ exp2 = tl.math.exp2 -+ log = tl.log -+ log2 = tl.log2 -+ -+ -+def get_all_max_shared_mem(): -+ try: -+ return [ -+ triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem'] -+ for i in range(device_torch_lib.device_count()) -+ ] -+ except BaseException: -+ _cpu_device_warning() -+ return [-1] -+ -+ -+class Backend(Enum): -+ ADA = 101376 # RTX 4090 -+ AMPERE = 166912 # A100 -+ HOPPER = 232448 # H100 -+ DEFAULT = 102400 # Default -+ -+ @classmethod -+ def get_shared_memory(cls, arch: str) -> int: -+ try: -+ return cls[arch.upper()].value -+ except KeyError: -+ return cls.DEFAULT.value -+ -+ -+@lru_cache(maxsize=None) -+def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: -+ try: -+ device_shared_mem_list = get_all_max_shared_mem() -+ max_shared_memory = device_shared_mem_list[tensor_idx] -+ return max_shared_memory >= Backend.get_shared_memory(arch) -+ except Exception: -+ return False -+ -+ -+def get_autotune_config( -+ multibuffer_list: tuple = (False,), -+ unit_flag_list: tuple = (False,), -+ limit_auto_multi_buffer_only_for_local_buffer_list: tuple = (False,), -+ limit_auto_multi_buffer_of_local_buffer_list: tuple = ("no-l0c",), -+ set_workspace_multibuffer_list: tuple = (2, 4), -+ enable_hivm_auto_cv_balance_list: tuple = (True,), -+ tile_mix_vector_loop_num_list: tuple = (2, 4), -+ tile_mix_cube_loop_num_list: tuple = (2, 4), -+): -+ configs = [] -+ for ( -+ multibuffer, -+ unit_flag, -+ limit_auto_multi_buffer_only_for_local_buffer, -+ limit_auto_multi_buffer_of_local_buffer, -+ ) in itertools.product( -+ list(multibuffer_list), -+ list(unit_flag_list), -+ list(limit_auto_multi_buffer_only_for_local_buffer_list), -+ list(limit_auto_multi_buffer_of_local_buffer_list), -+ ): -+ base_config_dict = { -+ 'multibuffer': multibuffer, -+ 'unit_flag': unit_flag, -+ 'limit_auto_multi_buffer_only_for_local_buffer': limit_auto_multi_buffer_only_for_local_buffer, -+ 'limit_auto_multi_buffer_of_local_buffer': limit_auto_multi_buffer_of_local_buffer, -+ } -+ -+ if limit_auto_multi_buffer_only_for_local_buffer: -+ configs.append(triton.Config(base_config_dict)) -+ else: -+ for ( -+ set_workspace_multibuffer, -+ enable_hivm_auto_cv_balance, -+ tile_mix_vector_loop, -+ tile_mix_cube_loop, -+ ) in itertools.product( -+ list(set_workspace_multibuffer_list), -+ list(enable_hivm_auto_cv_balance_list), -+ list(tile_mix_vector_loop_num_list), -+ list(tile_mix_cube_loop_num_list), -+ ): -+ full_config_dict = base_config_dict.copy() -+ full_config_dict.update( -+ { -+ 'set_workspace_multibuffer': set_workspace_multibuffer, -+ 'enable_hivm_auto_cv_balance': enable_hivm_auto_cv_balance, -+ 'tile_mix_vector_loop': tile_mix_vector_loop, -+ 'tile_mix_cube_loop': tile_mix_cube_loop, -+ } -+ ) -+ configs.append(triton.Config(full_config_dict)) -+ return configs -+ -+ -+def get_npu_properties(): -+ return driver.active.utils.get_device_properties(torch.npu.current_device()) -+ -+ -+@functools.cache -+def get_vector_num() -> int: -+ import torch_npu -+ -+ current_device = torch_npu.npu.current_device() -+ properties = driver.active.utils.get_device_properties(current_device) -+ return properties["num_vectorcore"] -+ -+ -+@lru_cache -+def is_arch35(): -+ try: -+ import torch_npu -+ -+ return "Ascend910_95" in torch_npu.npu.get_device_name() or "Ascend950" in torch_npu.npu.get_device_name() -+ except Exception: -+ return False -diff --git a/megatron/core/ssm/triton/wy_fast.py b/megatron/core/ssm/triton/wy_fast.py -new file mode 100644 -index 000000000..ab9c50582 ---- /dev/null -+++ b/megatron/core/ssm/triton/wy_fast.py -@@ -0,0 +1,340 @@ -+# -*- coding: utf-8 -*- -+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. -+ -+from typing import Optional, Tuple -+ -+import torch -+import triton -+import triton.language as tl -+ -+from .utils import prepare_chunk_indices, exp -+ -+ -+@triton.heuristics({ -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.jit(do_not_specialize=['T']) -+def prepare_wy_repr_bwd_kernel( -+ k, -+ v, -+ beta, -+ g, -+ A, -+ dw, -+ du, -+ dk, -+ dv, -+ dbeta, -+ dg, -+ cu_seqlens, -+ chunk_indices, -+ T, -+ B, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ NT: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ BV: tl.constexpr, -+ IS_VARLEN: tl.constexpr -+): -+ core_id = tl.program_id(0) -+ total_cores = tl.num_programs(0) -+ T_max = T -+ -+ base_chunks_per_pid = NT // total_cores -+ remainder_chunks = NT % total_cores -+ -+ if core_id < remainder_chunks: -+ chunks_this_pid = base_chunks_per_pid + 1 -+ start_idx = core_id * chunks_this_pid -+ else: -+ chunks_this_pid = base_chunks_per_pid -+ start_idx = core_id * chunks_this_pid + remainder_chunks -+ -+ for idx in range(start_idx, start_idx + chunks_this_pid): -+ for i_b in range(B): -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + idx * 2).to(tl.int32), tl.load(chunk_indices + idx * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ T = eos - bos -+ else: -+ i_t = idx -+ bos, eos = i_b * T, i_b * T + T -+ -+ o_t = i_t * BT + tl.arange(0, BT) -+ m_t = o_t < T -+ m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) -+ for i_h in range(0, H): -+ if IS_VARLEN: -+ offset = bos + i_h * T_max -+ else: -+ offset = bos * H + i_h * T_max -+ -+ p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) -+ -+ b_A = tl.load(p_A, boundary_check=(0, 1)) -+ b_beta = tl.load(p_beta, boundary_check=(0,)) -+ b_g = tl.load(p_g, boundary_check=(0,)) -+ b_g_exp = tl.exp(b_g) -+ -+ b_dbeta = tl.zeros([BT], dtype=tl.float32) -+ b_dA = tl.zeros([BT, BT], dtype=tl.float32) -+ b_dg = tl.zeros([BT], dtype=tl.float32) -+ -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ b_k_beta_g = (b_k * b_beta[:, None] * b_g_exp[:, None]).to(b_k.dtype) -+ b_dw = tl.load(p_dw, boundary_check=(0, 1)) -+ b_dA += tl.dot(b_dw, tl.trans(b_k_beta_g)) -+ b_dk_beta_g = tl.dot(b_A, b_dw) -+ b_dk = b_dk_beta_g * b_beta[:, None] * b_g_exp[:, None] -+ b_dbeta += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None], 1) -+ b_dg += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None] * b_beta[:, None], 1) -+ tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -+ -+ for i_v in range(tl.cdiv(V, BV)): -+ p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_du = tl.make_block_ptr(du + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ b_v = tl.load(p_v, boundary_check=(0, 1)) -+ b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) -+ b_du = tl.load(p_du, boundary_check=(0, 1)) -+ b_dA += tl.dot(b_du, tl.trans(b_v_beta)) -+ b_dv_beta = tl.dot(b_A, b_du) -+ b_dv = b_dv_beta * b_beta[:, None] -+ b_dbeta += tl.sum(b_dv_beta * b_v, 1) -+ tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) -+ -+ b_dA = tl.where(m_A, b_dA, 0) -+ b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) -+ b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) -+ b_dA = tl.where(m_A, -b_dA * exp(b_g[:, None] - b_g[None, :]), 0) -+ b_dA = b_dA.to(k.dtype.element_ty) -+ b_A = tl.zeros([BT, BT], dtype=tl.float32) -+ -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ b_dk = tl.load(p_dk, boundary_check=(0, 1)) -+ b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) -+ b_A += tl.dot(b_k_beta, tl.trans(b_k)) -+ b_dk_beta = tl.dot(b_dA, b_k) -+ b_dbeta += tl.sum(b_dk_beta * b_k, 1) -+ b_dk += tl.dot(tl.trans(b_dA), b_k_beta) -+ b_dk += b_dk_beta * b_beta[:, None] -+ tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) -+ -+ b_dA_A = b_dA * b_A -+ b_dg += tl.sum(b_dA_A, axis=1) - tl.sum(b_dA_A, axis=0) -+ p_dg = tl.make_block_ptr(dg + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ p_dbeta = tl.make_block_ptr(dbeta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) -+ tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) -+ -+ -+@triton.heuristics({ -+ 'USE_G': lambda args: args['g'] is not None, -+ 'USE_GK': lambda args: args['gk'] is not None, -+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -+}) -+@triton.jit(do_not_specialize=['T']) -+def recompute_w_u_fwd_kernel( -+ k, -+ v, -+ beta, -+ w, -+ u, -+ A, -+ g, -+ gk, -+ cu_seqlens, -+ chunk_indices, -+ T_tmp, -+ B, -+ H: tl.constexpr, -+ K: tl.constexpr, -+ V: tl.constexpr, -+ NT: tl.constexpr, -+ BT: tl.constexpr, -+ BK: tl.constexpr, -+ BV: tl.constexpr, -+ USE_G: tl.constexpr, -+ USE_GK: tl.constexpr, -+ IS_VARLEN: tl.constexpr -+): -+ core_id = tl.program_id(0) -+ total_cores = tl.num_programs(0) -+ T_max = T_tmp -+ -+ base_chunks_per_pid = NT // total_cores -+ remainder_chunks = NT % total_cores -+ -+ if core_id < remainder_chunks: -+ chunks_this_pid = base_chunks_per_pid + 1 -+ start_idx = core_id * chunks_this_pid -+ else: -+ chunks_this_pid = base_chunks_per_pid -+ start_idx = core_id * chunks_this_pid + remainder_chunks -+ -+ for idx in range(start_idx, start_idx + chunks_this_pid): -+ for i_b in range(B): -+ for i_h in range(0, H): -+ -+ if IS_VARLEN: -+ i_n, i_t = tl.load(chunk_indices + idx * 2).to(tl.int32), tl.load(chunk_indices + idx * 2 + 1).to(tl.int32) -+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) -+ offset = bos + i_h * T_max -+ T = eos - bos -+ else: -+ T = T_tmp -+ i_t = idx -+ bos, eos = i_b * T, i_b * T + T -+ offset = bos * H + i_h * T_max -+ -+ p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_beta = tl.load(p_beta, boundary_check=(0,)) -+ -+ p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) -+ b_A = tl.load(p_A, boundary_check=(0, 1)) -+ -+ for i_v in range(tl.cdiv(V, BV)): -+ p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) -+ b_v = tl.load(p_v, boundary_check=(0, 1)) -+ b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) -+ b_u = tl.dot(b_A, b_vb, allow_tf32=False) -+ tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) -+ -+ if USE_G: -+ p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,)) -+ b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) -+ -+ for i_k in range(tl.cdiv(K, BK)): -+ p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_k = tl.load(p_k, boundary_check=(0, 1)) -+ b_kb = b_k * b_beta[:, None] -+ if USE_G: -+ b_kb *= b_g[:, None] -+ if USE_GK: -+ p_gk = tl.make_block_ptr(gk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) -+ b_kb *= tl.exp(tl.load(p_gk, boundary_check=(0, 1))) -+ b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) -+ tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) -+ -+ -+def recompute_w_u_fwd( -+ k: torch.Tensor, -+ v: torch.Tensor, -+ beta: torch.Tensor, -+ A: torch.Tensor, -+ g: Optional[torch.Tensor] = None, -+ gk: Optional[torch.Tensor] = None, -+ cu_seqlens: Optional[torch.LongTensor] = None, -+) -> Tuple[torch.Tensor, torch.Tensor]: -+ B, T, H, K, V = *k.shape, v.shape[-1] -+ BT = A.shape[-1] -+ BK = 128 -+ BV = 128 -+ -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ g = g.transpose(1, 2).contiguous() if g is not None else None -+ beta = beta.transpose(1, 2).contiguous() -+ -+ w = torch.empty_like(k) -+ u = torch.empty_like(v) -+ cv_kernel_num = 24 -+ recompute_w_u_fwd_kernel[(cv_kernel_num,)]( -+ k=k, -+ v=v, -+ beta=beta, -+ w=w, -+ u=u, -+ A=A, -+ g=g, -+ gk=gk, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T_tmp=T, -+ B=B, -+ H=H, -+ K=K, -+ V=V, -+ NT=NT, -+ BT=BT, -+ BK=BK, -+ BV=BV, -+ ) -+ return w, u -+ -+ -+def prepare_wy_repr_bwd( -+ k: torch.Tensor, -+ v: torch.Tensor, -+ g: torch.Tensor, -+ beta: torch.Tensor, -+ A: torch.Tensor, -+ dw: torch.Tensor, -+ du: torch.Tensor, -+ cu_seqlens: Optional[torch.LongTensor], -+ chunk_size: int = 64, -+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: -+ B, T, H, K, V = *k.shape, v.shape[-1] -+ BT = chunk_size -+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None -+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) -+ BK = 128 -+ BV = 128 -+ beta = beta.transpose(1, 2).contiguous() -+ g = g.transpose(1, 2).contiguous() -+ -+ dk = torch.empty_like(k) -+ dv = torch.empty_like(v) -+ dbeta = torch.empty_like(beta) -+ dg = torch.empty_like(g) -+ -+ cv_kernel_num = 24 -+ prepare_wy_repr_bwd_kernel[(cv_kernel_num,)]( -+ k=k, -+ v=v, -+ beta=beta, -+ g=g, -+ A=A, -+ dw=dw, -+ du=du, -+ dk=dk, -+ dv=dv, -+ dbeta=dbeta, -+ dg=dg, -+ cu_seqlens=cu_seqlens, -+ chunk_indices=chunk_indices, -+ T=T, -+ B=B, -+ H=H, -+ K=K, -+ V=V, -+ NT=NT, -+ BT=BT, -+ BK=BK, -+ BV=BV, -+ ) -+ -+ dbeta = dbeta.transpose(1, 2).contiguous() -+ dg = dg.transpose(1, 2).contiguous() -+ -+ return dk, dv, dbeta, dg -+ -+ -+bwd_prepare_wy_repr = prepare_wy_repr_bwd -+ -+fwd_recompute_w_u = recompute_w_u_fwd +diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py +index 601a72a43..ec114cf38 100644 +--- a/megatron/core/ssm/gated_delta_net.py ++++ b/megatron/core/ssm/gated_delta_net.py +@@ -465,7 +465,7 @@ class GatedDeltaNet(MegatronModule): + + return out, out_bias + +- @jit_fuser ++ + def _apply_gated_norm(self, x, gate): + # Output Norm + x_dtype = x.dtype diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index bc5e4e2ee..d223e774e 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1219,7 +1219,7 @@ class Attention(MegatronModule, ABC): - + return output, bias - + - @jit_fuser -+ ++ def _apply_output_gate(self, x, gate): x_dtype = x.dtype gate = gate.contiguous() @@ -5365,15 +488,15 @@ index c30c107e7..f4340f4d1 100644 @@ -17,9 +17,9 @@ from megatron.core.transformer.utils import ( sharded_state_dict_default, ) - + -_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor) -_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor) -_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor) +_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor, torch.npu.FloatTensor) +_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor, torch.npu.HalfTensor) +_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor, torch.npu.BFloat16Tensor) - - + + def param_is_not_shared(param): # pylint: disable=missing-function-docstring diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index d8e753422..aff929985 100644 @@ -5382,18 +505,18 @@ index d8e753422..aff929985 100644 @@ -92,7 +92,7 @@ class GroupedMLP(MegatronModule): if self.config.activation_func not in (F.silu, F.gelu): raise ValueError("Activation function must be silu or gelu when using GroupedMLP.") - + - @jit_fuser -+ ++ def glu(x): x = torch.chunk(x, 2, dim=-1) return self.config.activation_func(x[0]) * x[1] @@ -109,7 +109,7 @@ class GroupedMLP(MegatronModule): "moe_act recompute for fp8 or fp4 cannot work with the legacy GroupedMLP." ) - + - @jit_fuser -+ ++ def activation_func_with_probs(x, probs): dtype = x.dtype res = self.activation_func(x) * probs @@ -5404,9 +527,9 @@ index 327dbc8a3..b007d6662 100644 @@ -1402,7 +1402,7 @@ class MoEFlexTokenDispatcher(MoETokenDispatcher): ).contiguous() return routing_map, probs - + - @jit_fuser -+ ++ def dispatch_preprocess( self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor ): @@ -5415,11 +538,11 @@ index 63f81465d..c47baeaea 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -7,6 +7,7 @@ from typing import Callable, List, Optional, Union - + import torch from torch import Tensor +import warnings - + from megatron.core import InferenceParams, parallel_state, tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedStateDict diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py @@ -5429,41 +552,20 @@ index d0ceca7af..f16796680 100644 @@ -69,7 +69,7 @@ class L2Norm(torch.nn.Module): self.hidden_size = hidden_size self.eps = eps - + - @jit_fuser -+ ++ def _norm(self, x): """ Performs the actual L2 normalization. -diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py -index 227e95862..da4540185 100644 ---- a/megatron/core/transformer/transformer_layer.py -+++ b/megatron/core/transformer/transformer_layer.py -@@ -775,6 +775,16 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): - self._set_fc2_residual(residual) - mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask) - -+ mlp_output, mlp_output_bias = mlp_output_with_bias -+ mlp_output = self.post_mlp_layernorm(mlp_output) -+ mlp_output_with_bias = (mlp_output, mlp_output_bias) -+ -+ if self.recompute_pre_mlp_layernorm: -+ # discard the output of the pre-mlp layernorm and register the recompute -+ # as a gradient hook of mlp_output_with_bias[0] -+ self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( -+ mlp_output_with_bias[0] -+ ) - nvtx_range_pop(suffix="mlp") - - if ( diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py index 880c53099..dbc95736e 100644 --- a/megatron/core/transformer/utils.py +++ b/megatron/core/transformer/utils.py @@ -51,7 +51,7 @@ def attention_mask_func(attention_scores, attention_mask): return attention_scores - - + + -@jit_fuser + def gelu_impl(x): @@ -5485,7 +587,7 @@ index e00e63148..ffe4b7ec6 100644 @@ -12,7 +12,7 @@ from megatron.core.jit import jit_fuser # actual gelu is: # x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) - + -@jit_fuser + def bias_gelu(bias, y): @@ -5506,8 +608,8 @@ index ca3414eec..5dd9e4798 100644 +++ b/megatron/legacy/model/transformer.py @@ -856,7 +856,7 @@ def get_bias_dropout_add(training): return _bias_dropout_add - - + + -@jit_fuser + def bias_dropout_add_fused_train(x: torch.Tensor, @@ -5515,8 +617,8 @@ index ca3414eec..5dd9e4798 100644 residual: torch.Tensor, @@ -864,7 +864,7 @@ def bias_dropout_add_fused_train(x: torch.Tensor, return bias_dropout_add(x, bias, residual, prob, True) - - + + -@jit_fuser + def bias_dropout_add_fused_inference(x: torch.Tensor, @@ -5528,32 +630,32 @@ index 5762000d5..534858df7 100644 +++ b/megatron/legacy/model/utils.py @@ -43,7 +43,7 @@ def get_linear_layer(rows, columns, init_method): return layer - - + + -@jit_fuser + def gelu_impl(x): """OpenAI's gelu implementation.""" return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * @@ -54,7 +54,7 @@ def openai_gelu(x): - - + + #This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter -@jit_fuser + def erf_gelu(x): return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype)+torch.ones_like(x).to(dtype=x.dtype)) - + diff --git a/megatron/training/utils.py b/megatron/training/utils.py index 7709f6513..65b8f7fca 100644 --- a/megatron/training/utils.py +++ b/megatron/training/utils.py @@ -446,6 +446,8 @@ def is_first_or_last_pipeline_stage(vp_stage): - + def get_device_arch_version(): """Returns GPU arch version (8: Ampere, 9: Hopper, 10: Blackwell, ...)""" + if hasattr(torch, 'npu') and torch.npu.is_available(): + return 10 # NPU: treat as Blackwell to avoid CUDA restrictions return torch.cuda.get_device_properties(torch.device("cuda:0")).major - - + + diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch index 0585af488..8a017c8e2 100644 --- a/docker/npu_patch/mindspeed.patch +++ b/docker/npu_patch/mindspeed.patch @@ -15,10 +15,23 @@ index bb007b44..98708a5b 100644 self.permute_idx_device = None input_chunk_idxs = torch.arange( diff --git a/mindspeed/core/megatron_basic/arguments_basic.py b/mindspeed/core/megatron_basic/arguments_basic.py -index 8ea25b9f..7853bce4 100644 +index 8ea25b9f..054313cb 100644 --- a/mindspeed/core/megatron_basic/arguments_basic.py +++ b/mindspeed/core/megatron_basic/arguments_basic.py -@@ -113,3 +113,35 @@ def transformer_config_init_wrapper(fn): +@@ -91,5 +91,11 @@ def transformer_config_init_wrapper(fn): + known_config = {} + unknown_config = {} + ignore_config = ['rope_type'] +- full_args = vars(get_full_args()).copy() ++ # vLLM serving arguments share the process-level Namespace in Vime, ++ # but they are not Megatron TransformerConfig fields. ++ full_args = { ++ key: value ++ for key, value in vars(get_full_args()).items() ++ if not key.startswith("vllm_") ++ } + full_args.update(dict(kwargs)) +@@ -113,3 +119,41 @@ def transformer_config_init_wrapper(fn): fn(self, *args, **known_config) return wrapper @@ -35,7 +48,13 @@ index 8ea25b9f..7853bce4 100644 +def transformer_config_init_subclass(cls, **kwargs): + mutable_types = (list, dict, set, bytearray) + unknown_config = {} -+ full_args = vars(get_full_args()).copy() ++ # Keep serving-only vLLM values out of dynamically created Megatron ++ # dataclass fields. They remain available on the original Vime args. ++ full_args = { ++ key: value ++ for key, value in vars(get_full_args()).items() ++ if not key.startswith("vllm_") ++ } + full_args.update(kwargs) + + config_key = inspect.signature(cls).parameters diff --git a/docker/npu_patch/series.conf b/docker/npu_patch/series.conf index 55c054365..10d543dee 100644 --- a/docker/npu_patch/series.conf +++ b/docker/npu_patch/series.conf @@ -11,6 +11,7 @@ # Ordinary patch changes must not add patch-specific logic to the CI script. /vllm-workspace/vllm|vllm.patch|docker/npu_patch/vllm.patch /vllm-workspace/vllm-ascend|vllm-ascend.patch|docker/npu_patch/vllm-ascend.patch +/root/Megatron-LM|megatron-common.patch|docker/patch/latest/megatron.patch /root/Megatron-LM|megatron.patch|docker/npu_patch/megatron.patch /root/Megatron-Bridge|megatron-bridge.patch|docker/npu_patch/megatron-bridge.patch /root/MindSpeed|mindspeed.patch|docker/npu_patch/mindspeed.patch diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index fa89a9129..7dd2da745 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -1,53 +1,472 @@ +diff --git a/vllm_ascend/distributed/weight_transfer/__init__.py b/vllm_ascend/distributed/weight_transfer/__init__.py +index d6434f05f..fd7223836 100644 +--- a/vllm_ascend/distributed/weight_transfer/__init__.py ++++ b/vllm_ascend/distributed/weight_transfer/__init__.py +@@ -33,6 +33,11 @@ def register_engine(): + "vllm_ascend.distributed.weight_transfer.npu_ipc_engine", + "NPUIPCWeightTransferEngine", + ) ++ WeightTransferTrainerFactory.register_engine( ++ "hccl", ++ "vllm_ascend.distributed.weight_transfer.hccl_engine", ++ "HCCLTrainerWeightTransferEngine", ++ ) + WeightTransferTrainerFactory.register_engine( + "npu_ipc", + "vllm_ascend.distributed.weight_transfer.npu_ipc_engine", +diff --git a/vllm_ascend/distributed/weight_transfer/hccl_engine.py b/vllm_ascend/distributed/weight_transfer/hccl_engine.py +index 023ffd1ae..7b184e279 100644 +--- a/vllm_ascend/distributed/weight_transfer/hccl_engine.py ++++ b/vllm_ascend/distributed/weight_transfer/hccl_engine.py +@@ -3,8 +3,9 @@ + """HCCL-based weight transfer engine.""" + + from collections.abc import Callable, Iterator +-from dataclasses import dataclass +-from typing import TYPE_CHECKING, Any ++from concurrent.futures import ThreadPoolExecutor ++from dataclasses import asdict, dataclass ++from typing import TYPE_CHECKING, Any, ClassVar + + import torch + +@@ -14,6 +15,11 @@ if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.base import ( ++ ParamMeta, ++ TrainerInitInfo, ++ TrainerWeightTransferEngine, ++ VLLMWeightSyncClient, ++ WeightSource, + WeightTransferEngine, + WeightTransferInitInfo, + WeightTransferUpdateInfo, +@@ -26,6 +32,19 @@ from vllm_ascend.distributed.weight_transfer.packed_tensor import ( + ) + + ++@dataclass ++class HCCLTrainerInitInfo(TrainerInitInfo): ++ """Stateful trainer configuration; rank 0 owns the HCCL endpoint.""" ++ ++ backend: ClassVar[str] = "hccl" ++ master_address: str ++ master_port: int ++ world_size: int ++ packed: bool = True ++ packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES ++ packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS ++ ++ + @dataclass + class HCCLWeightTransferInitInfo(WeightTransferInitInfo): + """Initialization info for HCCL weight transfer backend.""" +@@ -338,3 +357,114 @@ class HCCLWeightTransferEngine(WeightTransferEngine[HCCLWeightTransferInitInfo, + pg = StatelessProcessGroup.create(host=master_address, port=master_port, rank=rank, world_size=world_size) + pyhccl = PyHcclCommunicator(pg, device=device) + return pyhccl ++ ++ ++class HCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[HCCLTrainerInitInfo]): ++ """Stateful control plane over the existing HCCL broadcast transport. ++ ++ Every trainer rank replays the source collectives. Only rank 0 opens the ++ transfer communicator and drives worker RPCs. The HCCL receiver expects ++ packed geometry on each update, so derive it from the trainer init info. ++ """ ++ ++ init_info_cls = HCCLTrainerInitInfo ++ ++ def __init__(self, *, init_info: HCCLTrainerInitInfo, client: VLLMWeightSyncClient, source: WeightSource) -> None: ++ super().__init__(client=client, source=source, is_sender=init_info.is_sender) ++ self.init_info = init_info ++ self.model_update_group: PyHcclCommunicator | None = None ++ ++ @classmethod ++ def trainer_init( ++ cls, ++ init_info: HCCLTrainerInitInfo, ++ *, ++ client: VLLMWeightSyncClient, ++ source: WeightSource, ++ ) -> "HCCLTrainerWeightTransferEngine": ++ engine = cls(init_info=init_info, client=client, source=source) ++ if not engine.is_sender: ++ return engine ++ ++ worker_info = HCCLWeightTransferInitInfo( ++ master_address=init_info.master_address, ++ master_port=init_info.master_port, ++ rank_offset=1, ++ world_size=init_info.world_size, ++ ) ++ # Both endpoints must rendezvous concurrently. ++ executor = ThreadPoolExecutor(max_workers=1) ++ try: ++ future = executor.submit(client.init_weight_transfer_engine, asdict(worker_info)) ++ if future.done(): ++ future.result() ++ engine.model_update_group = HCCLWeightTransferEngine.trainer_init(worker_info) ++ future.result() ++ finally: ++ executor.shutdown(wait=False) ++ return engine ++ ++ def send_weights(self) -> None: ++ # Metadata export can itself contain Megatron collectives. ++ meta = self.source.metadata() ++ if not self.is_sender: ++ for _ in self.source: ++ pass ++ torch.npu.current_stream().synchronize() ++ return ++ ++ assert self.model_update_group is not None, "HCCL trainer has been shut down." ++ info = self.init_info ++ update_info = HCCLWeightTransferUpdateInfo( ++ names=[m.name for m in meta], ++ dtype_names=[str(m.dtype).split(".")[-1] for m in meta], ++ shapes=[list(m.shape) for m in meta], ++ packed=info.packed, ++ packed_buffer_size_bytes=info.packed_buffer_size_bytes, ++ packed_num_buffers=info.packed_num_buffers, ++ ) ++ self.client.start_weight_update() ++ executor = ThreadPoolExecutor(max_workers=1) ++ try: ++ future = executor.submit(self.client.update_weights, asdict(update_info)) ++ if future.done(): ++ future.result() ++ HCCLWeightTransferEngine.trainer_send_weights( ++ self._checked_iter(self.source, meta), ++ HCCLTrainerSendWeightsArgs( ++ group=self.model_update_group, ++ packed=info.packed, ++ packed_buffer_size_bytes=info.packed_buffer_size_bytes, ++ packed_num_buffers=info.packed_num_buffers, ++ ), ++ ) ++ future.result() ++ finally: ++ # A failed broadcast can leave the worker RPC waiting in HCCL. ++ # Surface the error; joining that thread here would hide it forever. ++ executor.shutdown(wait=False) ++ self.client.finish_weight_update() ++ torch.npu.current_stream().synchronize() ++ ++ @staticmethod ++ def _checked_iter(source: WeightSource, meta: list[ParamMeta]) -> Iterator[tuple[str, torch.Tensor]]: ++ # Receive buffer sizes and packed chunk boundaries come from metadata. ++ sent = 0 ++ for name, tensor in source: ++ if sent >= len(meta): ++ raise ValueError(f"WeightSource yielded more parameters than metadata(): {name!r}.") ++ expected = meta[sent] ++ if name != expected.name or tensor.dtype != expected.dtype or tuple(tensor.shape) != expected.shape: ++ raise ValueError( ++ f"WeightSource metadata() disagrees with iteration at index {sent}: " ++ f"expected {expected.name!r} {expected.dtype} {expected.shape}, " ++ f"got {name!r} {tensor.dtype} {tuple(tensor.shape)}." ++ ) ++ sent += 1 ++ # Unpacked HCCL reads contiguous storage from data_ptr(). ++ yield name, tensor if tensor.is_contiguous() else tensor.contiguous() ++ if sent != len(meta): ++ raise ValueError(f"WeightSource yielded {sent} parameters but metadata() declared {len(meta)}.") ++ ++ def shutdown(self) -> None: ++ self.model_update_group = None +diff --git a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py +index 7a63adf7b..d3e3a7fc4 100644 +--- a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py ++++ b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py +@@ -156,12 +156,14 @@ class NPUIPCWeightTransferEngine( # type: ignore[no-redef] + self.packed = init_info.packed + + def start_weight_update(self) -> None: +- """No-op for NPU IPC engine (no layerwise reloading).""" +- pass ++ from vllm.model_executor.model_loader.reload import initialize_layerwise_reload ++ ++ initialize_layerwise_reload(self.model) + + def finish_weight_update(self) -> None: +- """No-op for NPU IPC engine (no layerwise reloading).""" +- pass ++ from vllm.model_executor.model_loader.reload import finalize_layerwise_reload ++ ++ finalize_layerwise_reload(self.model, self.model_config) + + def receive_weights(self, update_info: NPUIPCWeightTransferUpdateInfo) -> None: + """Receive weights from the trainer via NPU IPC handles. +@@ -170,6 +172,10 @@ class NPUIPCWeightTransferEngine( # type: ignore[no-redef] + update_info: NPU IPC update info containing parameter names, + dtypes, shapes, and IPC handles. + """ ++ # A rank-local slot may have no weights in this chunk (and no handle). ++ if not update_info.names: ++ return ++ + # Use the worker's assigned device rather than the ambient current + # device: the receive path is no longer wrapped in + # ``with torch.device(self.device)`` by the caller, so the current +@@ -219,7 +225,7 @@ class NPUIPCWeightTransferEngine( # type: ignore[no-redef] + weight = rebuild_npu_tensor(*list_args) + weights.append((name, weight)) + +- self.model.load_weights(weights) ++ self.model.load_weights(weights) + + def shutdown(self) -> None: + pass +@@ -288,6 +294,8 @@ class NPUIPCTrainerWeightTransferEngine(IPCTrainerWeightTransferEngine): + self.client.finish_weight_update() + self._post_send_sync() + del weight_refs ++ torch.npu.ipc_collect() ++ torch.npu.empty_cache() + + def _send(self, source: "WeightSource") -> list[torch.Tensor] | None: + if self.packed: +diff --git a/vllm_ascend/distributed/weight_transfer/packed_tensor.py b/vllm_ascend/distributed/weight_transfer/packed_tensor.py +index a35d9af8d..d55c988c3 100644 +--- a/vllm_ascend/distributed/weight_transfer/packed_tensor.py ++++ b/vllm_ascend/distributed/weight_transfer/packed_tensor.py +@@ -39,6 +39,7 @@ def packed_broadcast_producer( + target_packed_tensor_size = buffer_size_bytes + + streams = [torch.npu.Stream() for _ in range(num_buffers)] ++ source_stream = torch.npu.current_stream() + buffer_idx = 0 + + packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)] +@@ -50,6 +51,8 @@ def packed_broadcast_producer( + # Synchronize the current stream (waits for previous + # iteration's work on this buffer to finish) + streams[buffer_idx].synchronize() ++ # Wait for tensors produced on the caller's stream before packing. ++ streams[buffer_idx].wait_stream(source_stream) + # Start tasks for the new buffer in a new stream + with torch.npu.stream(streams[buffer_idx]): + # Initialize the packing tensor list and sizes +@@ -206,10 +209,10 @@ def packed_npu_ipc_producer( + ) -> Iterator[dict[str, Any]]: + """Pack tensors into a reusable NPU IPC buffer and yield chunks. + +- Allocates a single NPU buffer of ``buffer_size_bytes`` and registers +- it for IPC once via ``reduce_tensor``. Each chunk's packed data is +- copied into this buffer before yielding, so only one IPC-shared +- allocation exists for the lifetime of the transfer. ++ Allocates a single NPU buffer of ``buffer_size_bytes``. Each chunk ++ publishes a fresh IPC reference via ``reduce_tensor`` so its consumer ++ releases that reference exactly once. The underlying buffer is reused ++ for the lifetime of the transfer. + + Args: + iterator: Iterator of (name, tensor) pairs. +@@ -218,9 +221,6 @@ def packed_npu_ipc_producer( + buffer_size_bytes: Exact capacity of the reusable IPC buffer. + """ + ipc_buffer = torch.empty(buffer_size_bytes, dtype=torch.uint8, device="npu") +- # Store only the rebuild args (drop the func); the consumer rebuilds with +- # the well-known ``rebuild_npu_tensor``, mirroring upstream's CUDA IPC engine. +- _, ipc_args = reduce_tensor(ipc_buffer) + + names: list[str] = [] + shapes: list[list[int]] = [] +@@ -240,6 +240,7 @@ def packed_npu_ipc_producer( + + if total_bytes and total_bytes + flat.numel() > buffer_size_bytes: + torch.npu.current_stream().synchronize() ++ _, ipc_args = reduce_tensor(ipc_buffer) + yield { + "names": names, + "shapes": shapes, +@@ -259,6 +260,7 @@ def packed_npu_ipc_producer( + + if total_bytes: + torch.npu.current_stream().synchronize() ++ _, ipc_args = reduce_tensor(ipc_buffer) + yield { + "names": names, + "shapes": shapes, +diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py +index c388619f2..38a7d2fb7 100644 +--- a/vllm_ascend/worker/v2/model_runner.py ++++ b/vllm_ascend/worker/v2/model_runner.py +@@ -17,7 +17,7 @@ + # This file is a part of the vllm-ascend project. + # + +-from contextlib import contextmanager ++from contextlib import AbstractContextManager, contextmanager + + import numpy as np + import torch +@@ -239,9 +239,13 @@ class NPUModelRunner(GPUModelRunner): + self.pp_handler.broadcast_draft_tokens() + return output + +- def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: ++ def initialize_kv_cache( ++ self, ++ kv_cache_config: KVCacheConfig, ++ kv_cache_allocation_context: AbstractContextManager | None = None, ++ ) -> None: + with graph_manager_wrapper(self): +- super().initialize_kv_cache(kv_cache_config) ++ super().initialize_kv_cache(kv_cache_config, kv_cache_allocation_context=kv_cache_allocation_context) + if self.pcp_manager is not None: + assert isinstance(self.pcp_manager, AscendPCPManager) + self.pcp_manager.vllm_config = self.vllm_config diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py -index 062b3ecd..dd87affb 100644 +index 6d99ac76a..b7df4df22 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py -@@ -404,7 +404,7 @@ class NPUWorker(WorkerBase): +@@ -321,7 +321,48 @@ class NPUWorker(WorkerBase): + + def start_weight_update(self) -> None: + """Begin a new weight update; prepares the model for layerwise reload.""" ++ with set_current_vllm_config(self.vllm_config): ++ self._start_weight_update() ++ ++ def get_draft_model(self) -> nn.Module | None: ++ return self.model_runner.get_draft_model() ++ ++ def supports_draft_weight_updates(self) -> bool: ++ engine = self.weight_transfer_engine ++ speculative_config = self.speculative_config ++ get_draft_model = getattr(self.model_runner, "get_draft_model", None) ++ return ( ++ engine is not None ++ and engine.supports_draft_weight_update ++ and callable(get_draft_model) ++ and get_draft_model() is not None ++ and speculative_config is not None ++ and speculative_config.draft_model_config is not None ++ ) ++ ++ def _set_draft_weight_update_target(self) -> None: ++ assert self.weight_transfer_engine is not None ++ draft_model = self.get_draft_model() ++ if draft_model is None: ++ raise RuntimeError("Draft model weight update requested, but no draft model is configured.") ++ speculative_config = self.speculative_config ++ if speculative_config is None or speculative_config.draft_model_config is None: ++ raise RuntimeError("Draft model weight update requested, but no draft model config is configured.") ++ self.weight_transfer_engine.set_weight_update_target(draft_model, speculative_config.draft_model_config) ++ ++ def start_draft_weight_update(self) -> None: ++ """Retarget the native engine at the draft model for this session.""" ++ with set_current_vllm_config(self.vllm_config): ++ self._start_weight_update(is_draft=True) ++ ++ def _start_weight_update(self, is_draft: bool = False) -> None: + self._check_weight_transfer_engine() ++ assert self.weight_transfer_engine is not None ++ ++ if is_draft and not self.weight_transfer_engine.supports_draft_weight_update: ++ raise RuntimeError( ++ f"{type(self.weight_transfer_engine).__name__} does not support draft model weight updates." ++ ) + + if self._weight_update_active: + raise RuntimeError( +@@ -330,11 +371,16 @@ class NPUWorker(WorkerBase): + + self._check_nz_disabled() + +- assert self.weight_transfer_engine is not None +- self.weight_transfer_engine.start_weight_update() ++ try: ++ if is_draft: ++ self._set_draft_weight_update_target() ++ self.weight_transfer_engine.start_weight_update() ++ except BaseException: ++ self.weight_transfer_engine.reset_weight_update_target() ++ raise + self._weight_update_active = True + +- def update_weights(self, update_info: dict) -> None: ++ def update_weights(self, update_info: dict | list[dict]) -> None: + """Receive a chunk of weights from the trainer and load them in place.""" + self._check_weight_transfer_engine() + assert self.weight_transfer_engine is not None +@@ -343,11 +389,19 @@ class NPUWorker(WorkerBase): + if not self._weight_update_active: + raise RuntimeError("start_weight_update must be called before update_weights.") + +- try: +- self.weight_transfer_engine.update_weights(update_info) +- except BaseException: +- self._weight_update_active = False +- raise ++ with set_current_vllm_config(self.vllm_config): ++ try: ++ if isinstance(update_info, list): ++ parallel_config = self.vllm_config.parallel_config ++ worker_rank = parallel_config.data_parallel_rank * parallel_config.world_size + self.rank ++ local_update_info = update_info[worker_rank] ++ else: ++ local_update_info = update_info ++ self.weight_transfer_engine.update_weights(local_update_info) ++ except BaseException: ++ self._weight_update_active = False ++ self.weight_transfer_engine.reset_weight_update_target() ++ raise + + def finish_weight_update(self) -> None: + """Finish the current weight update; runs layerwise postprocessing.""" +@@ -357,8 +411,12 @@ class NPUWorker(WorkerBase): + raise RuntimeError("start_weight_update must be called before finish_weight_update.") + + assert self.weight_transfer_engine is not None +- self.weight_transfer_engine.finish_weight_update() +- self._weight_update_active = False ++ with set_current_vllm_config(self.vllm_config): ++ try: ++ self.weight_transfer_engine.finish_weight_update() ++ finally: ++ self._weight_update_active = False ++ self.weight_transfer_engine.reset_weight_update_target() + + def shutdown(self) -> None: + if ensure_kv_transfer_shutdown is not None: +@@ -448,7 +506,9 @@ class NPUWorker(WorkerBase): # take current memory snapshot - self.init_snapshot = MemorySnapshot() + self.init_snapshot = MemorySnapshot(device=device) self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization - if self.init_snapshot.free_memory < self.requested_memory: -+ if False and self.init_snapshot.free_memory < self.requested_memory: # colocate ++ weight_transfer_config = self.vllm_config.weight_transfer_config ++ uses_ipc_weight_transfer = weight_transfer_config is not None and weight_transfer_config.backend == "npu_ipc" ++ if not uses_ipc_weight_transfer and self.init_snapshot.free_memory < self.requested_memory: GiB = lambda b: round(b / GiB_bytes, 2) raise ValueError( f"Free memory on device " -@@ -535,15 +535,16 @@ class NPUWorker(WorkerBase): - self.npugraph_memory_estimate = npugraph_memory_estimate - +@@ -597,7 +657,9 @@ class NPUWorker(WorkerBase): + self.non_torch_memory = profile_result.non_torch_increase + free_gpu_memory = profile_result.after_profile.free_memory - assert self.init_snapshot.free_memory > free_gpu_memory, ( -- "Error in memory profiling. " -- f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, " -- f"current free memory {GiB(free_gpu_memory)} GiB. " -- "This happens when other processes sharing the same container " -- "release GPU memory while vLLM is profiling during initialization. " -- "To fix this, ensure consistent GPU memory allocation or " -- "isolate vLLM in its own container." -- ) -+ # colocate: skip memory profiling assert -+ # assert self.init_snapshot.free_memory > free_gpu_memory, ( -+ # "Error in memory profiling. " -+ # f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, " -+ # f"current free memory {GiB(free_gpu_memory)} GiB. " -+ # "This happens when other processes sharing the same container " -+ # "release GPU memory while vLLM is profiling during initialization. " -+ # "To fix this, ensure consistent GPU memory allocation or " -+ # "isolate vLLM in its own container." -+ # ) - self.available_kv_cache_memory_bytes = ( - self.requested_memory - profile_result.non_kv_cache_memory - npugraph_memory_estimate_applied - ) -diff --git a/vllm_ascend/distributed/weight_transfer/packed_tensor.py b/vllm_ascend/distributed/weight_transfer/packed_tensor.py -index a35d9af8d..66a179cd6 100644 ---- a/vllm_ascend/distributed/weight_transfer/packed_tensor.py -+++ b/vllm_ascend/distributed/weight_transfer/packed_tensor.py -@@ -41,2 +41,3 @@ def packed_broadcast_producer( - streams = [torch.npu.Stream() for _ in range(num_buffers)] -+ source_stream = torch.npu.current_stream() - buffer_idx = 0 -@@ -52,2 +53,5 @@ def packed_broadcast_producer( - streams[buffer_idx].synchronize() -+ # Source tensors may have been produced asynchronously on the caller's -+ # stream. Wait before the packing stream reads them in torch.cat. -+ streams[buffer_idx].wait_stream(source_stream) - # Start tasks for the new buffer in a new stream ++ weight_transfer_config = self.vllm_config.weight_transfer_config ++ uses_ipc_weight_transfer = weight_transfer_config is not None and weight_transfer_config.backend == "npu_ipc" ++ assert uses_ipc_weight_transfer or self.init_snapshot.free_memory > free_gpu_memory, ( + "Error in memory profiling. " + f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, " + f"current free memory {GiB(free_gpu_memory)} GiB. " +@@ -1118,8 +1180,12 @@ class NPUWorker(WorkerBase): + from contextlib import nullcontext + + context = nullcontext() # type: ignore +- with context: +- self.model_runner.initialize_kv_cache(kv_cache_config) ++ if self.use_v2_model_runner: ++ # MRV2 bookkeeping must survive sleep; pool only the KV data. ++ self.model_runner.initialize_kv_cache(kv_cache_config, kv_cache_allocation_context=context) ++ else: ++ with context: ++ self.model_runner.initialize_kv_cache(kv_cache_config) + + # MRV2's scheduler emits new_block_ids_to_zero whenever this flag is + # set, so its worker-side consumer must use the same condition. Keep the diff --git a/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch index 976fea2b9..142941525 100644 --- a/docker/npu_patch/vllm.patch +++ b/docker/npu_patch/vllm.patch @@ -1,42 +1,158 @@ -diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py -index 2e407ae..5e4c5de 100644 ---- a/vllm/model_executor/layers/rotary_embedding/common.py -+++ b/vllm/model_executor/layers/rotary_embedding/common.py -@@ -135,7 +135,7 @@ class ApplyRotaryEmb(CustomOp): - self.enable_fp32_compute = enable_fp32_compute - - self.apply_rotary_emb_flash_attn = None -- if not current_platform.is_cpu() and find_spec("flash_attn") is not None: -+ if not current_platform.is_cpu() and find_spec("flash_attn") is not None and not hasattr(current_platform, "is_npu"): - from flash_attn.ops.triton.rotary import apply_rotary - - self.apply_rotary_emb_flash_attn = apply_rotary -diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py -index cb61bca..5c076d5 100644 ---- a/vllm/v1/core/sched/async_scheduler.py -+++ b/vllm/v1/core/sched/async_scheduler.py -@@ -51,7 +51,7 @@ class AsyncScheduler(Scheduler): - - # Update the number of output placeholders. - request.num_output_placeholders -= len(new_token_ids) -- assert request.num_output_placeholders >= 0 -+ request.num_output_placeholders = max(0, request.num_output_placeholders) - - # Cache the new tokens. Preempted requests should be skipped. - if status_before_update == RequestStatus.RUNNING: +diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +index f304bf677b..71a17e9363 100644 +--- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py ++++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +@@ -240,6 +240,8 @@ class GenerateStreamResponse(BaseModel): + ) + choices: list[GenerateResponseStreamChoice] + usage: UsageInfo | None = Field(default=None) ++ weight_version: str | None = None ++ request_spec_decode_stats: dict[str, Any] | None = Field(default=None) + + + class GenerateResponse(BaseModel): +@@ -255,6 +257,8 @@ class GenerateResponse(BaseModel): + created: int | None = None + choices: list[GenerateResponseChoice] + usage: UsageInfo | None = Field(default=None) ++ weight_version: str | None = None ++ request_spec_decode_stats: dict[str, Any] | None = Field(default=None) + prompt_logprobs: list[dict[int, Logprob] | None] | None = None + + kv_transfer_params: dict[str, Any] | None = Field( +diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py +index bbbd85137c..809f1a66ea 100644 +--- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py ++++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py +@@ -14,6 +14,7 @@ from vllm.engine.protocol import EngineClient + from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker + from vllm.entrypoints.generate.base.serving import ( + GenerateBaseServing, ++ build_spec_decoding_metrics, + clamp_prompt_logprobs, + ) + from vllm.entrypoints.openai.chat_completion.protocol import ( +@@ -261,6 +262,7 @@ class ServingTokens(GenerateBaseServing): + ) + + assert result_generator is not None ++ weight_version = await self.engine_client.get_weight_version() + + if request.stream: + return self.serve_tokens_stream_generator( +@@ -269,10 +271,16 @@ class ServingTokens(GenerateBaseServing): + request_id, + model_name, + request_metadata, ++ weight_version, + ) + + return await self.serve_tokens_full_generator( +- request, result_generator, request_id, model_name, request_metadata ++ request, ++ result_generator, ++ request_id, ++ model_name, ++ request_metadata, ++ weight_version, + ) + + async def serve_tokens_full_generator( +@@ -282,6 +290,7 @@ class ServingTokens(GenerateBaseServing): + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, ++ weight_version: str | None, + ) -> ErrorResponse | GenerateResponse: + created_time = int(time.time()) + final_res: RequestOutput | None = None +@@ -355,6 +364,11 @@ class ServingTokens(GenerateBaseServing): + cached_tokens=final_res.num_cached_tokens + ) + ++ spec_decode_metrics = build_spec_decoding_metrics(final_res) ++ request_spec_decode_stats = ( ++ spec_decode_metrics.model_dump() if spec_decode_metrics else None ++ ) ++ + request_metadata.final_usage_info = usage + + response = GenerateResponse( +@@ -363,6 +377,8 @@ class ServingTokens(GenerateBaseServing): + model=model_name, + choices=choices, + usage=usage, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), + kv_transfer_params=final_res.kv_transfer_params, + ec_transfer_params=final_res.ec_transfer_params, +@@ -396,11 +412,13 @@ class ServingTokens(GenerateBaseServing): + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, ++ weight_version: str | None, + ) -> AsyncGenerator[str, None]: + num_prompt_tokens = 0 + num_generated_tokens: list[int] = [] + first_iteration = True + num_cached_tokens = None ++ request_spec_decode_stats: dict[str, object] | None = None + sampling_params: SamplingParams = request.sampling_params + + include_usage, include_continuous_usage = should_include_usage( +@@ -409,6 +427,9 @@ class ServingTokens(GenerateBaseServing): + + try: + async for res in result_generator: ++ spec_decode_metrics = build_spec_decoding_metrics(res) ++ if spec_decode_metrics is not None: ++ request_spec_decode_stats = spec_decode_metrics.model_dump() + if first_iteration: + if res.prompt_token_ids is not None: + num_prompt_tokens = len(res.prompt_token_ids) +@@ -448,6 +469,8 @@ class ServingTokens(GenerateBaseServing): + + chunk = GenerateStreamResponse( + request_id=request_id, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + choices=[ + GenerateResponseStreamChoice( + index=i, +@@ -482,6 +505,8 @@ class ServingTokens(GenerateBaseServing): + if include_usage: + final_chunk = GenerateStreamResponse( + request_id=request_id, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + choices=[], + usage=final_usage_info, + ) +diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py +index a8f8023cf2..d62f39a2b8 100644 +--- a/vllm/model_executor/model_loader/reload/meta.py ++++ b/vllm/model_executor/model_loader/reload/meta.py +@@ -30,5 +30,6 @@ SKIP_LOAD_TENSORS: set[str] = { + "expert_global_to_physical", + "expert_physical_to_global", + "expert_local_to_global", ++ "expert_ids_per_ep_rank", + "e_score_correction_bias", + } diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py -index 596cb48..8c30495 100644 +index 1d30a0eaf6..37239f7f62 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py -@@ -126,7 +126,10 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module): +@@ -129,6 +129,10 @@ class Glm4MoeLiteMultiTokenPredictorLayer(nn.Module): ) -> torch.Tensor: assert inputs_embeds is not None # masking inputs at position 0, as not needed by MTP - inputs_embeds[positions == 0] = 0 -+ # torch.where (element-wise) replaces bool-mask index to stay -+ # cudagraph-capturable on NPU (aclnnNonzeroV2 fails under capture). ++ # Avoid dynamic bool indexing during NPU graph capture. + mask = (positions == 0).unsqueeze(-1) -+ inputs_embeds = torch.where(mask, torch.zeros_like(inputs_embeds), inputs_embeds) ++ inputs_embeds = torch.where( ++ mask, torch.zeros_like(inputs_embeds), inputs_embeds ++ ) inputs_embeds = self.enorm(inputs_embeds) previous_hidden_states = self.hnorm(previous_hidden_states) - diff --git a/docker/patch/latest/vllm-pull_weights.patch b/docker/patch/latest/vllm-pull_weights.patch new file mode 100644 index 000000000..ddd67eb49 --- /dev/null +++ b/docker/patch/latest/vllm-pull_weights.patch @@ -0,0 +1,505 @@ +diff --git a/tests/model_executor/model_loader/test_local_checkpoint.py b/tests/model_executor/model_loader/test_local_checkpoint.py +new file mode 100644 +index 0000000000..192b7e979a +--- /dev/null ++++ b/tests/model_executor/model_loader/test_local_checkpoint.py +@@ -0,0 +1,144 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import json ++import zlib ++ ++import numpy as np ++import pytest ++import safetensors.numpy ++import zstandard ++ ++from vllm.utils.local_checkpoint import pull_checkpoint ++ ++ ++def _checksum(data: np.ndarray) -> str: ++ return f"{zlib.adler32(data):08x}" ++ ++ ++def _write_delta( ++ source_dir, ++ version: int, ++ old: np.ndarray, ++ new: np.ndarray, ++ encoding: str, ++) -> None: ++ version_dir = source_dir / f"weight_v{version:06d}" ++ version_dir.mkdir() ++ old_bytes = old.view(np.uint8).reshape(-1) ++ new_bytes = new.view(np.uint8).reshape(-1) ++ if encoding == "xor": ++ payload = old_bytes ^ new_bytes ++ else: ++ positions = np.flatnonzero(old_bytes != new_bytes).astype(" None: ++ version_dir = source_dir / f"weight_v{version:06d}" ++ version_dir.mkdir() ++ safetensors.numpy.save_file({"weight": weight}, version_dir / "model.safetensors") ++ (version_dir / "config.json").write_text("{}") ++ ++ ++@pytest.mark.parametrize("encoding", ["xor", "overwrite"]) ++def test_pull_checkpoint_applies_vime_delta(tmp_path, encoding): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ updated = baseline.copy() ++ updated[0, 1] = 100.0 ++ updated[2, 3] = -5.0 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, updated, encoding) ++ ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 0) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ ++ actual = safetensors.numpy.load_file(local_dir / "model.safetensors") ++ np.testing.assert_array_equal(actual["weight"], updated) ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000001"} ++ ++ ++def test_pull_checkpoint_resets_to_latest_full_version(tmp_path): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ first = baseline + 1 ++ reset = baseline + 10 ++ latest = reset.copy() ++ latest[1, 2] = -7.0 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, first, "xor") ++ _write_full(source_dir, 2, reset) ++ _write_delta(source_dir, 3, reset, latest, "xor") ++ ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 3) ++ ++ actual = safetensors.numpy.load_file(local_dir / "model.safetensors") ++ np.testing.assert_array_equal(actual["weight"], latest) ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000003"} ++ ++ ++def test_pull_checkpoint_does_not_advance_on_checksum_failure(tmp_path): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ updated = baseline + 1 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, updated, "xor") ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 0) ++ safetensors.numpy.save_file( ++ {"weight": baseline + 2}, local_dir / "model.safetensors" ++ ) ++ ++ with pytest.raises(RuntimeError, match="Checksum mismatch"): ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000000"} +diff --git a/vllm/utils/local_checkpoint.py b/vllm/utils/local_checkpoint.py +new file mode 100644 +index 0000000000..42be58d249 +--- /dev/null ++++ b/vllm/utils/local_checkpoint.py +@@ -0,0 +1,319 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Maintain a host-local checkpoint from full and delta weight versions.""" ++ ++from __future__ import annotations ++ ++import fcntl ++import glob ++import importlib ++import io ++import json ++import mmap ++import os ++import shutil ++import struct ++import threading ++import zlib ++from concurrent.futures import ThreadPoolExecutor ++from contextlib import ExitStack, contextmanager, suppress ++ ++import numpy as np ++import zstandard ++ ++NUM_WORKERS = min(32, os.cpu_count() or 8) ++SYNC_DIR = ".weight_sync" ++ ++ ++def pull_checkpoint( ++ local_checkpoint_dir: str, ++ base_dir: str, ++ source_dir: str, ++ target_version: int, ++ pre_read_hook: str | None = None, ++) -> None: ++ """Bring a host-local checkpoint to a published weight version.""" ++ if target_version > 0 and pre_read_hook: ++ module_path, _, function_name = pre_read_hook.rpartition(".") ++ hook = getattr(importlib.import_module(module_path), function_name) ++ hook(source_dir, target_version) ++ with _pull_lock(local_checkpoint_dir): ++ applied = _read_applied_version(local_checkpoint_dir) ++ floor = applied if applied is not None else 0 ++ start = target_version ++ while start > floor and _is_delta(_version_dir(source_dir, start)): ++ start -= 1 ++ ++ if applied is None or start > applied: ++ seed_dir = base_dir if start == 0 else _version_dir(source_dir, start) ++ _reset_checkpoint(seed_dir, local_checkpoint_dir, start) ++ else: ++ start = applied ++ ++ for version in range(start + 1, target_version + 1): ++ _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version)) ++ ++ ++def _version_dir(source_dir: str, version: int) -> str: ++ return os.path.join(source_dir, f"weight_v{version:06d}") ++ ++ ++def _is_delta(version_dir: str) -> bool: ++ if not os.path.isdir(version_dir): ++ raise FileNotFoundError(f"Published weight version missing: {version_dir}") ++ try: ++ with open( ++ os.path.join(version_dir, "model.safetensors.index.json") ++ ) as index_file: ++ return "delta_encoding" in json.load(index_file).get("metadata", {}) ++ except FileNotFoundError: ++ return False ++ ++ ++class _Adler32: ++ def __init__(self) -> None: ++ self._value = 1 ++ ++ def update(self, data) -> None: ++ self._value = zlib.adler32(data, self._value) ++ ++ def hexdigest(self) -> str: ++ return f"{self._value:08x}" ++ ++ ++def _new_hasher(algorithm: str): ++ if algorithm == "xxh3-128": ++ import xxhash ++ ++ return xxhash.xxh3_128() ++ if algorithm == "blake3": ++ import blake3 ++ ++ return blake3.blake3() ++ if algorithm == "adler32": ++ return _Adler32() ++ raise KeyError(f"Unknown checksum algorithm {algorithm!r}") ++ ++ ++def _checksum(algorithm: str, data) -> str: ++ hasher = _new_hasher(algorithm) ++ hasher.update(data) ++ return hasher.hexdigest() ++ ++ ++@contextmanager ++def _pull_lock(local_checkpoint_dir: str): ++ sync_dir = os.path.join(local_checkpoint_dir, SYNC_DIR) ++ os.makedirs(sync_dir, exist_ok=True) ++ with open(os.path.join(sync_dir, "lock"), "w") as lock_file: ++ fcntl.flock(lock_file, fcntl.LOCK_EX) ++ try: ++ yield ++ finally: ++ fcntl.flock(lock_file, fcntl.LOCK_UN) ++ ++ ++def _read_applied_version(local_checkpoint_dir: str) -> int | None: ++ try: ++ with open( ++ os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json") ++ ) as state_file: ++ return int(json.load(state_file)["version"]) ++ except FileNotFoundError: ++ return None ++ ++ ++def _write_applied_version(local_checkpoint_dir: str, version: int) -> None: ++ path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json") ++ temporary = f"{path}.tmp" ++ with open(temporary, "w") as state_file: ++ json.dump({"version": f"{version:06d}"}, state_file) ++ state_file.flush() ++ os.fsync(state_file.fileno()) ++ os.replace(temporary, path) ++ ++ ++def _drop_page_cache(path: str) -> None: ++ try: ++ file_descriptor = os.open(path, os.O_RDONLY) ++ try: ++ os.posix_fadvise(file_descriptor, 0, 0, os.POSIX_FADV_DONTNEED) ++ finally: ++ os.close(file_descriptor) ++ except OSError: ++ pass ++ ++ ++def _reset_checkpoint(source_dir: str, local_checkpoint_dir: str, version: int) -> None: ++ os.makedirs(local_checkpoint_dir, exist_ok=True) ++ source_files = [entry for entry in os.scandir(source_dir) if entry.is_file()] ++ for entry in source_files: ++ shutil.copy2(entry.path, os.path.join(local_checkpoint_dir, entry.name)) ++ _drop_page_cache(entry.path) ++ ++ source_names = {entry.name for entry in source_files} ++ for entry in os.scandir(local_checkpoint_dir): ++ if entry.is_file() and entry.name not in source_names: ++ os.remove(entry.path) ++ ++ for entry in source_files: ++ copied_size = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name)) ++ if copied_size != entry.stat().st_size: ++ raise RuntimeError( ++ f"Size mismatch copying {entry.name}: " ++ f"source {entry.stat().st_size} != local {copied_size}" ++ ) ++ _write_applied_version(local_checkpoint_dir, version) ++ ++ ++def _tensor_locations(checkpoint_dir: str) -> dict[str, tuple[str, int, int]]: ++ locations = {} ++ for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")): ++ with open(path, "rb") as tensor_file: ++ (header_length,) = struct.unpack(" None: ++ with open(os.path.join(version_dir, "model.safetensors.index.json")) as index_file: ++ metadata = json.load(index_file)["metadata"] ++ ++ applied = _read_applied_version(local_checkpoint_dir) ++ version = int(metadata["version"]) ++ if applied == version: ++ return ++ if applied != int(metadata["base_version"]): ++ raise RuntimeError( ++ f"Out-of-order delta: local at {applied}, " ++ f"delta builds on {metadata['base_version']}" ++ ) ++ if metadata["compression_format"] != "zstd": ++ raise NotImplementedError( ++ f"Compression {metadata['compression_format']!r} is not supported" ++ ) ++ ++ encoding = metadata["delta_encoding"] ++ checksum_algorithm = metadata["checksum_format"] ++ locations = _tensor_locations(local_checkpoint_dir) ++ open_mmaps = {} ++ resources = ExitStack() ++ mismatches = [] ++ mismatch_lock = threading.Lock() ++ delta_blobs = [] ++ items = [] ++ try: ++ for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): ++ with open(delta_file, "rb") as tensor_file: ++ blob = tensor_file.read() ++ delta_blobs.append(blob) ++ (header_length,) = struct.unpack(" None: ++ with mismatch_lock: ++ mismatches.append(name) ++ ++ def apply_xor(item) -> None: ++ name, compressed, path, offset, byte_count, expected = item ++ region = np.ndarray( ++ (byte_count,), ++ dtype=np.uint8, ++ buffer=open_mmaps[path], ++ offset=offset, ++ ) ++ hasher = _new_hasher(checksum_algorithm) ++ reader = zstandard.ZstdDecompressor().stream_reader( ++ io.BytesIO(bytes(compressed)) ++ ) ++ position = 0 ++ while position < byte_count: ++ block = reader.read(min(2 << 20, byte_count - position)) ++ if not block: ++ break ++ chunk = np.frombuffer(block, dtype=np.uint8) ++ region[position : position + chunk.size] ^= chunk ++ hasher.update(region[position : position + chunk.size]) ++ position += chunk.size ++ if position != byte_count or hasher.hexdigest() != expected: ++ report_mismatch(name) ++ ++ def apply_overwrite(item) -> None: ++ name, compressed, path, offset, byte_count, expected = item ++ delta = np.frombuffer( ++ zstandard.ZstdDecompressor().decompress(bytes(compressed)), ++ dtype=np.uint8, ++ ) ++ region = np.ndarray( ++ (byte_count,), ++ dtype=np.uint8, ++ buffer=open_mmaps[path], ++ offset=offset, ++ ) ++ count = int.from_bytes(delta[:4].tobytes(), "little") ++ positions_end = 4 + 4 * count ++ positions = np.frombuffer(delta[4:positions_end].tobytes(), dtype=" dict[str, Any]: ++ from vllm.utils.local_checkpoint import pull_checkpoint ++ ++ pull_checkpoint( ++ local_checkpoint_dir=local_checkpoint_dir, ++ base_dir=self.model_config.model, ++ source_dir=source_dir, ++ target_version=target_version, ++ pre_read_hook=pre_read_hook, ++ ) ++ ++ return {"success": True, "weight_version": str(target_version)} ++ + @torch.inference_mode() + def determine_available_memory(self) -> int: + """Profiles the peak memory usage of the model to determine how much diff --git a/docker/patch/latest/vllm.patch b/docker/patch/latest/vllm.patch index e278cf3f0..580932836 100644 --- a/docker/patch/latest/vllm.patch +++ b/docker/patch/latest/vllm.patch @@ -1,30 +1,168 @@ -diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py ---- a/vllm/v1/engine/core.py -+++ b/vllm/v1/engine/core.py -@@ -750,8 +750,10 @@ - if tags is None or tags: - self.model_executor.wake_up(tags) - -- # Resume scheduling (applies to all levels) -- self.resume_scheduler() -+ # Partial wakes intentionally keep the remaining allocations asleep. -+ # Resume scheduling only once all executor memory is resident again. -+ if not self.model_executor.is_sleeping: -+ self.resume_scheduler() - - def is_sleeping(self) -> bool: - """Check if engine is sleeping at any level.""" -@@ -1844,8 +1846,11 @@ - continue - - # We are in a running state and so must execute a dummy pass -- # if the model didn't execute any ready requests. -- self.execute_dummy_batch() -+ # if the model didn't execute any ready requests -- unless the executor is -+ # asleep (#44483: a decode-shaped dummy batch reads freed KV -> illegal memory -+ # access). The finished-sync all-reduce below still runs (DP lockstep). -+ if not self.is_sleeping(): -+ self.execute_dummy_batch() - - # 3) All-reduce operation to determine global unfinished reqs. - self.engines_running = self._has_global_unfinished_reqs( \ No newline at end of file +diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +index f304bf677ba..71a17e9363d 100644 +--- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py ++++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +@@ -240,6 +240,8 @@ class GenerateStreamResponse(BaseModel): + ) + choices: list[GenerateResponseStreamChoice] + usage: UsageInfo | None = Field(default=None) ++ weight_version: str | None = None ++ request_spec_decode_stats: dict[str, Any] | None = Field(default=None) + + + class GenerateResponse(BaseModel): +@@ -255,6 +257,8 @@ class GenerateResponse(BaseModel): + created: int | None = None + choices: list[GenerateResponseChoice] + usage: UsageInfo | None = Field(default=None) ++ weight_version: str | None = None ++ request_spec_decode_stats: dict[str, Any] | None = Field(default=None) + prompt_logprobs: list[dict[int, Logprob] | None] | None = None + + kv_transfer_params: dict[str, Any] | None = Field( +diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py +index bbbd85137ce..809f1a66ea5 100644 +--- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py ++++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py +@@ -14,6 +14,7 @@ from vllm.engine.protocol import EngineClient + from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker + from vllm.entrypoints.generate.base.serving import ( + GenerateBaseServing, ++ build_spec_decoding_metrics, + clamp_prompt_logprobs, + ) + from vllm.entrypoints.openai.chat_completion.protocol import ( +@@ -261,6 +262,7 @@ class ServingTokens(GenerateBaseServing): + ) + + assert result_generator is not None ++ weight_version = await self.engine_client.get_weight_version() + + if request.stream: + return self.serve_tokens_stream_generator( +@@ -269,10 +271,16 @@ class ServingTokens(GenerateBaseServing): + request_id, + model_name, + request_metadata, ++ weight_version, + ) + + return await self.serve_tokens_full_generator( +- request, result_generator, request_id, model_name, request_metadata ++ request, ++ result_generator, ++ request_id, ++ model_name, ++ request_metadata, ++ weight_version, + ) + + async def serve_tokens_full_generator( +@@ -282,6 +290,7 @@ class ServingTokens(GenerateBaseServing): + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, ++ weight_version: str | None, + ) -> ErrorResponse | GenerateResponse: + created_time = int(time.time()) + final_res: RequestOutput | None = None +@@ -355,6 +364,11 @@ class ServingTokens(GenerateBaseServing): + cached_tokens=final_res.num_cached_tokens + ) + ++ spec_decode_metrics = build_spec_decoding_metrics(final_res) ++ request_spec_decode_stats = ( ++ spec_decode_metrics.model_dump() if spec_decode_metrics else None ++ ) ++ + request_metadata.final_usage_info = usage + + response = GenerateResponse( +@@ -363,6 +377,8 @@ class ServingTokens(GenerateBaseServing): + model=model_name, + choices=choices, + usage=usage, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), + kv_transfer_params=final_res.kv_transfer_params, + ec_transfer_params=final_res.ec_transfer_params, +@@ -396,11 +412,13 @@ class ServingTokens(GenerateBaseServing): + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, ++ weight_version: str | None, + ) -> AsyncGenerator[str, None]: + num_prompt_tokens = 0 + num_generated_tokens: list[int] = [] + first_iteration = True + num_cached_tokens = None ++ request_spec_decode_stats: dict[str, object] | None = None + sampling_params: SamplingParams = request.sampling_params + + include_usage, include_continuous_usage = should_include_usage( +@@ -409,6 +427,9 @@ class ServingTokens(GenerateBaseServing): + + try: + async for res in result_generator: ++ spec_decode_metrics = build_spec_decoding_metrics(res) ++ if spec_decode_metrics is not None: ++ request_spec_decode_stats = spec_decode_metrics.model_dump() + if first_iteration: + if res.prompt_token_ids is not None: + num_prompt_tokens = len(res.prompt_token_ids) +@@ -448,6 +469,8 @@ class ServingTokens(GenerateBaseServing): + + chunk = GenerateStreamResponse( + request_id=request_id, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + choices=[ + GenerateResponseStreamChoice( + index=i, +@@ -482,6 +505,8 @@ class ServingTokens(GenerateBaseServing): + if include_usage: + final_chunk = GenerateStreamResponse( + request_id=request_id, ++ weight_version=weight_version, ++ request_spec_decode_stats=request_spec_decode_stats, + choices=[], + usage=final_usage_info, + ) +diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py +--- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py ++++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py +@@ -45,6 +45,6 @@ from vllm.compilation.decorators import support_torch_compile + from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig + from vllm.config.speech_to_text import SpeechToTextParams +-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size ++from vllm.distributed import get_pp_group + from vllm.inputs import PromptType + from vllm.logger import init_logger + from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY +@@ -188,8 +188,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module): + self.embed_dim = config.d_model + self.num_heads = config.encoder_attention_heads + self.head_dim = self.embed_dim // self.num_heads +- tp_size = get_tensor_model_parallel_world_size() +- self.num_local_heads = self.num_heads // tp_size ++ self.num_local_heads = self.num_heads + + if (self.head_dim * self.num_heads) != self.embed_dim: + raise ValueError( +@@ -213,6 +212,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module): + total_num_kv_heads=self.num_heads, + bias=True, + prefix=f"{prefix}.qkv_proj", ++ disable_tp=True, + ) + + self.out_proj = RowParallelLinear( +@@ -213,6 +213,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module): + output_size=self.embed_dim, + bias=True, + prefix=f"{prefix}.out_proj", ++ disable_tp=True, + ) + + self.attn = MMEncoderAttention( diff --git a/docker/version.txt b/docker/version.txt index 5a83f4bc4..2b2bb2c29 100644 --- a/docker/version.txt +++ b/docker/version.txt @@ -1 +1 @@ -nightly-dev-20260519a +nightly-dev-20260828a diff --git a/docs/_static/image/logo.ico b/docs/_static/image/logo.ico index 78153d1e2..751fd934a 100644 Binary files a/docs/_static/image/logo.ico and b/docs/_static/image/logo.ico differ diff --git a/docs/_static/image/logo.jpg b/docs/_static/image/logo.jpg index c1257d72b..22db26759 100644 Binary files a/docs/_static/image/logo.jpg and b/docs/_static/image/logo.jpg differ diff --git a/docs/conf.py b/docs/conf.py index 90e62460d..0e1fcd83c 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -4,9 +4,9 @@ from datetime import datetime from pathlib import Path -sys.path.insert(0, os.path.abspath("../..")) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -__version__ = "0.0.1" +__version__ = "0.3.2" project = "Vime" copyright = f"2025-{datetime.now().year}, Vime" @@ -185,7 +185,7 @@ def _sync_examples(app): if not candidate.exists(): continue # skip entirely if nothing suitable target_dir = out_dir / d.name - target_dir.mkdir(parents=True, exist_ok=True) + shutil.copytree(d, target_dir, ignore=shutil.ignore_patterns("README.md", "README_zh.md")) shutil.copy2(candidate, target_dir / "README.md") entries.append((d.name, f"_examples_synced/{d.name}/README.md")) diff --git a/docs/en/advanced/arch-support-beyond-megatron.md b/docs/en/advanced/arch-support-beyond-megatron.md index e3b5335d6..aca455cc9 100644 --- a/docs/en/advanced/arch-support-beyond-megatron.md +++ b/docs/en/advanced/arch-support-beyond-megatron.md @@ -22,8 +22,8 @@ vime leverages this mechanism by **hijacking the spec generation stage to replac * **Corresponding File**: `vime_plugins/models/hf_attention.py` 3. **Aligning Model Weights** - Once the model architecture is integrated, we must ensure that the weights can be loaded correctly. We use the [mbridge](https://github.com/ISEEKYAN/mbridge) library, through our `Qwen3NextBridge`, to establish a naming map between the HuggingFace checkpoint and Megatron's parameters, enabling seamless, bidirectional conversion. - * **Corresponding File**: `vime_plugins/mbridge/qwen3_next.py` + Once the model architecture is integrated, we must ensure that the weights can be loaded correctly. vime keeps the HuggingFace-to-Megatron name mapping and tensor transforms next to its checkpoint loader. + * **Corresponding File**: `vime/backends/megatron_utils/hf_to_megatron/qwen3_next.py` Through the coordination of these three components, we can successfully run a complex model architecture not natively supported by Megatron—using its HuggingFace implementation as the vehicle—on top of Megatron's parallel framework. This is achieved while fully retaining all key capabilities like model parallelism, MoE acceleration, and pipeline scheduling. diff --git a/docs/en/advanced/delta-weight-sync.md b/docs/en/advanced/delta-weight-sync.md new file mode 100644 index 000000000..41d1a95d2 --- /dev/null +++ b/docs/en/advanced/delta-weight-sync.md @@ -0,0 +1,104 @@ +# Delta Weight Sync + +Delta weight sync keeps non-colocated rollout engines up to date by shipping only the bytes +that changed between two syncs, instead of a full checkpoint each time. It targets large-model +training/inference disaggregation across clusters or datacenters, where writing the whole actor +every sync is the dominant cost. + +It is **disk-transport only**. The trainer publishes each sync as a canonical HF checkpoint +directory; the engine's `/pull_weights` endpoint (shipped in vime's vllm patch) fans the +apply out to **every host the engine spans** and verifies it, then the engine reloads the +patched local checkpoint through the **ordinary** `update_weights_from_disk` endpoint. vime +only ever talks to one endpoint per engine, so multi-node serving and external rollout engines +need nothing extra on the vime side. + +## Configuration + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/delta-updates +--update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt +--update-weight-delta-encoding xor # or: overwrite +--update-weight-delta-checksum xxh3-128 # or: blake3, adler32 +``` + +| Flag | Role | +|---|---| +| `--update-weight-disk-dir` | Shared filesystem directory the trainer publishes deltas to and the rollout hosts read from. | +| `--update-weight-local-checkpoint-dir` | Host-local (e.g. NVMe) full HF checkpoint that `/pull_weights` keeps in sync — deltas are applied into it in place; a published full checkpoint replaces it. Each host seeds it from the engine's model path on the first `/pull_weights`. | +| `--update-weight-delta-encoding` | On-disk delta encoding: `xor` (default) or `overwrite`. | +| `--update-weight-delta-checksum` | Per-tensor integrity checksum: `xxh3-128` (default), `blake3`, or `adler32`. | + +Deltas are always zstd-compressed (level 1); profiling showed it dominates lz4 / gzip / snappy / brotli on both wire size and decompress speed for this data, so it is not a knob. + +## How it works + +1. **Seed.** On the first sync the trainer captures a CPU snapshot of every parameter — seeded + from `--hf-checkpoint`, which is exactly what each rollout host materializes its local + checkpoint from. Nothing is published; this snapshot is the base the next sync diffs against. + The trainer also issues `/pull_weights` with `target_version=0` so every host materializes + its local base now, overlapped with the snapshot capture. +2. **Publish.** On every later sync the trainer diffs each gathered HF tensor against the + snapshot, encodes and compresses the change, and writes a new version directory + `weight_v{N:06d}/` under `--update-weight-disk-dir`. The directory is a canonical HF + checkpoint — `model-NNNNN.safetensors` files holding the compressed diff tensors plus a + `model.safetensors.index.json` (tensor name → file) carrying the apply metadata — so the + artifact is portable, not tied to the trainer's parallelism layout. The snapshot is then + advanced to the new values for the next diff. +3. **Pull.** The trainer calls `/pull_weights` on each engine. Inside the engine the request is + broadcast to every rank on every node; each host applies the new version's delta into its + local checkpoint in place (a per-host file lock collapses co-located ranks to one apply). + The apply is parallelized across tensors and verified per-tensor (see Integrity); the call + only reports success once **every host** holds a checksum-verified checkpoint. + + `/pull_weights` is not delta-specific: each published version is self-describing, and a + version that is an ordinary full HF checkpoint (no delta metadata in its index) is pulled by + copying it as-is — resetting the chain, so a fresh host joining late seeds from the newest + full version instead of replaying every delta, and older deltas can be pruned. vime's + full-mode disk sync uses exactly this when `--update-weight-local-checkpoint-dir` is set. +4. **Reload.** The engines reload the patched local checkpoint through the vanilla + `update_weights_from_disk` path — the weight-loading code never sees the delta format. + +Because the snapshot is seeded from `--hf-checkpoint` (the engine's actual base) rather than +from the current GPU weights, the scheme is correct for any model even where the Megatron→HF +round-trip is not byte-exact (e.g. trimmed vocab-padding rows in the embedding / LM head). + +## Encodings + +Both encodings are byte-level and dtype-blind, so the same path works for quantized checkpoints. +The engine reads the choice from each version's index metadata. + +- **`xor`** (default): writes `new ^ old`. Smallest wire and fastest to apply (sequential, + cache-friendly; the unchanged bytes are zeros the compressor crushes). It is an involution, + so it must be applied **exactly once** against the correct base — applying it twice reverts. +- **`overwrite`**: writes the changed positions and their new absolute values. Larger on the + wire and a less cache-friendly scattered apply, but **idempotent**: re-applying it (or + finishing a partially-applied delta) converges to the same state regardless of how many times + it runs. Use it when re-applicability matters more than wire size. + +## Integrity + +The trainer stores a per-tensor checksum of each tensor's new state in the version. After +applying, every host recomputes the checksum and **raises on any mismatch** — the failure +propagates through the `/pull_weights` response, so a corrupt delta or a wrong base fails loud +instead of serving bad weights. The apply also refuses to run out of order: a version only +applies on top of its declared base version. + +`--update-weight-delta-checksum` selects the algorithm. The checksum is not the apply bottleneck +(the apply is decompress + XOR bound), so this is a digest-property choice, not a speed one: +`xxh3-128` (default) is the widest fast non-cryptographic digest; `blake3` is cryptographic, for +untrusted storage; `adler32` is for interop with systems that expect it. + +## Shared-filesystem visibility hooks + +On a POSIX shared filesystem (NFS, Lustre, …) no extra step is needed. Object-store-backed +mounts that need an explicit publish/refresh to make writes visible across hosts can supply two +optional hooks, loaded by import path — no vendor-specific code lives in vime or vllm: + +- `--custom-update-weight-post-write-path` (vime, trainer side): called after a version's files are + written, before the engines are told to read it (e.g. upload pending writes to the backing object store). + Signature: `hook(args, version_dir, rollout_engines)`. +- `--custom-update-weight-pre-read-path` (vime, engine side): called on each host + inside the engine before `/pull_weights` reads the delta directory (e.g. refresh the mount's view). + Signature: `hook(delta_dir, target_version)`. diff --git a/docs/en/advanced/external-rollout-engines.md b/docs/en/advanced/external-rollout-engines.md new file mode 100644 index 000000000..f9fb82cb2 --- /dev/null +++ b/docs/en/advanced/external-rollout-engines.md @@ -0,0 +1,110 @@ +# External Rollout Engines Roadmap + +An external rollout engine is a vLLM engine that is not launched by the vime training job. Another system deploys and owns the engine lifecycle; vime connects to those engines during training, registers a router, and syncs updated actor weights when needed. + +This page is a roadmap. Use it to decide when to use `--rollout-external-engine-addrs`, when to stay with `--vllm-config`, and which weight-update path to pick for external deployments. + +## Where To Start + +| Goal | Recommended entry point | +| :--- | :--- | +| Engines are already launched externally and vime should only connect for rollout | `--rollout-external-engine-addrs` | +| vime should still launch engines, but you need PD disaggregation, multi-model serving, heterogeneous server groups, or per-group overrides | [vLLM Config](vllm-config.md) | +| Trainer and external engines can form an NCCL group | Default `--update-weight-mode full --update-weight-transport nccl` | +| Trainer and external engines cannot form an NCCL group, but can see the same filesystem path | `--update-weight-mode full --update-weight-transport disk` | +| Full checkpoints are too heavy for large-model cross-cluster or cross-DC sync | `--update-weight-mode delta --update-weight-transport disk` | +| Rollout serving can use an independent vLLM environment, or even different GPU models/vendors | external engines + disk transport | +| You need frozen reference, reward, or tool-side models | Prefer `update_weights: false` in [vLLM Config](vllm-config.md#3-multi-model-serving) | + +Delta mode supports disk transport only. Use full mode when syncing weights over NCCL. + +## What External Engine Does + +First launch vLLM servers independently: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10090 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10091 ... +``` + +Then pass those addresses to the training job: + +```bash +python train.py \ + --rollout-external-engine-addrs host1:10090 host2:10091 \ + ... +``` + +vime queries each engine's `/server_info` or `/get_server_info` endpoint and infers GPU counts, TP/PP information, and worker type (`regular`, `prefill`, or `decode`). If `--vllm-router-ip/--vllm-router-port` is not provided, vime launches its own router and registers the external engines with it. + +This path fits deployments where serving is owned outside the training job: a separate inference cluster, a separate Ray cluster, manually warmed vLLM engines, or a rollout service managed by another orchestrator. + +## Relationship With `--vllm-config` + +`--rollout-external-engine-addrs` and `--vllm-config` are mutually exclusive because they own different boundaries: + +- `--vllm-config`: vime owns the engine lifecycle. The YAML describes the topology, and vime launches server groups, routers, multi-model serving, and selective weight updates. +- `--rollout-external-engine-addrs`: an external system owns the engine lifecycle. vime discovers already-running engines, attaches them to a router, and treats them as the default rollout model. + +If your main requirement is multi-model serving, frozen reference/reward models, PD disaggregation, or heterogeneous group configuration, prefer `--vllm-config`. Use external engines when the engines are already deployed outside the training job. + +## Environment And Hardware Decoupling + +An important implication of external engines is that the vLLM serving side does not need to use the vime training job's Python environment, Megatron environment, or Ray runtime. It can run in a separate vLLM container, an independent cluster, or another orchestration system. vime only depends on the HTTP endpoint, `/server_info`, and the communication path required by the selected weight-sync transport. + +With disk transport, weights move through HF checkpoints or safetensors deltas on a shared filesystem, and vLLM hot-loads them through `update_weights_from_disk`. This path does not require the training GPUs and rollout GPUs to be the same model, or even from the same vendor, as long as vLLM supports that hardware backend, model format, and precision configuration. For example, training can run on one GPU fleet while rollout serving runs on another fleet with different GPU models or vendors. + +With NCCL transport, the usual NCCL communication and hardware-compatibility requirements still apply. For cross-vendor, incompatible-network, or cross-datacenter deployments, prefer `--update-weight-transport disk`. + +## Update From Disk + +Full-checkpoint update from disk is the simplest fallback path for external deployments: + +```bash +--update-weight-mode full +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/full-updates +``` + +At every weight sync, the trainer writes a complete HF checkpoint directory under `--update-weight-disk-dir`, such as `weight_v000123/`, then calls each vLLM engine's `update_weights_from_disk` endpoint over HTTP so the engine reloads the checkpoint without a process restart. + +Adding `--update-weight-local-checkpoint-dir` makes each engine first pull the published checkpoint onto every host it spans (`/pull_weights`, shipped in vime's vllm patch) and reload from local disk (e.g. NVMe) — one shared-filesystem read per host instead of one per rank, which matters when the shared dir is object-store-backed or the engine spans several nodes. + +This mode has a simple control plane: it does not require an NCCL group between trainer and engines. It only requires both sides to see the same shared filesystem path. The tradeoff is size: every sync writes the full actor weights, which is expensive for large models or frequent updates. + +For debugging, add: + +```bash +--update-weight-disk-keep-files +``` + +This keeps the full-checkpoint directories after engines acknowledge the load. + +## Update With Delta + +Delta update targets large-model training/inference disaggregation across clusters or datacenters. Instead of writing a full checkpoint every sync, the trainer keeps a CPU snapshot of the previous sync, diffs each parameter against it, and publishes only the changed bytes; each engine's `/pull_weights` endpoint (shipped in vime's vllm patch) applies the delta into a host-local checkpoint on every host the engine spans, and the engine reloads via the vanilla `update_weights_from_disk` endpoint. vime only calls the engine's HTTP endpoint, so multi-node external engines work the same as vime-launched ones. + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/delta-updates +--update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt +``` + +See [Delta Weight Sync](delta-weight-sync.md) for the mechanism, encodings, integrity checks, and shared-filesystem visibility hooks. + +## Deployment Checklist + +- External engine HTTP addresses must be reachable from the training job. +- External engines can use an independent vLLM environment; they do not need the vime or Megatron training environment. +- Disk transport supports different GPU models or vendors between training and rollout, as long as vLLM supports the target hardware and model format. +- Disk transport requires trainer and vLLM engines to see the same `--update-weight-disk-dir` path; a path visible only to the trainer is not enough. +- External engines are not recovered by vime fault tolerance; their lifecycle belongs to the external deployment system. +- `--vllm-config` and `--rollout-external-engine-addrs` are mutually exclusive. +- Delta mode does not support `--colocate`, because colocated sync uses CUDA IPC handles and delta encoding does not reduce the actual transfer. + +## Related Work + +The [Composer 2 technical report by the Cursor Research Team](https://arxiv.org/html/2603.24477v2) describes a similar production shape: training and rollout generation run asynchronously, Cursor partners with Fireworks AI for RL inference, updated weights are written to shared S3 every training step, delta compression reduces transfer size, and inference clusters in different regions download and reconstruct weights from a shared delta chain. + +vime's external engines, update from disk, and delta disk transport address the same infrastructure problem: once training and inference are disaggregated, weight sync must work across processes, clusters, and even datacenters without letting full-model transfer dominate the training loop. diff --git a/docs/en/advanced/fault-tolerance.md b/docs/en/advanced/fault-tolerance.md index f4767c812..97e504bcf 100644 --- a/docs/en/advanced/fault-tolerance.md +++ b/docs/en/advanced/fault-tolerance.md @@ -12,8 +12,8 @@ Enable fault tolerance with: vime currently provides rollout-engine fault tolerance: -- health checks for vLLM rollout engines; -- timeout-based rollout engine restart; +- health checks for vLLM rollout servers; +- timeout-based rollout server restart; - correct parameter update after restart; - debug rollout dumps for replaying training-side issues without rerunning rollout; - trace/profiling hooks for inspecting long-tail rollout behavior. @@ -22,7 +22,7 @@ Cluster-level preemption, trainer-rank failure, and full-job resume should still ## Rollout Health Checks -During rollout, vime periodically sends heartbeat requests (`/health`) to all vLLM engines. If a heartbeat times out, the unhealthy vLLM engine is stopped. After the current rollout round completes, vime restarts the engine and updates it with the correct parameters before it serves future rollout requests. +During rollout, vime periodically sends heartbeat requests (`/health`) to all vLLM servers. If a heartbeat times out, the unhealthy vLLM server is stopped. After the current rollout round completes, vime restarts the server and updates it with the correct parameters before it serves future rollout requests. The main arguments are: @@ -65,7 +65,7 @@ For long-running jobs: - If startup health checks fail on large MoE models, increase `--rollout-health-check-first-wait`. - If transient load spikes cause false positives, increase `--rollout-health-check-timeout`. -- If an engine repeatedly restarts after weight sync, inspect the vLLM logs and the latest rollout debug dump. +- If a server repeatedly restarts after weight sync, inspect the vLLM logs and the latest rollout debug dump. - If the trainer fails rather than rollout, resume from checkpoint and use debug replay to isolate whether the saved rollout batch is valid. ## Related Docs diff --git a/docs/en/advanced/low-precision.md b/docs/en/advanced/low-precision.md new file mode 100644 index 000000000..e91ca0d6a --- /dev/null +++ b/docs/en/advanced/low-precision.md @@ -0,0 +1,143 @@ +# Low Precision Training and Rollout + +Low precision in vime is primarily used to make rollout faster and more memory-efficient while keeping training numerically stable. For large MoE RL jobs, the recommended production path is: + +> **BF16 training in Megatron + FP8 rollout/inference in vLLM** + +Megatron keeps the trainable checkpoint in BF16/torch_dist format. vLLM serves an FP8 Hugging Face checkpoint for rollout. During weight updates, vime uses the quantization config in `--hf-checkpoint` to quantize updated BF16 weights before sending them to vLLM. + +## Feature Maturity + +| Feature | Status | Recommended Use | +|---|---|---| +| BF16 training + FP8 rollout/inference | Stable | Default path for large MoE RL recipes. Keeps training stable while reducing rollout memory and bandwidth. | +| FP8 KV cache in vLLM rollout | Stable when supported by your vLLM version/GPU stack | Increase KV cache capacity for long-context or agentic rollout by passing `--vllm-kv-cache-dtype fp8_e4m3`. | +| INT4 rollout / INT4 QAT | Beta | Use when rollout memory/throughput pressure is high and the model path has been validated. | +| FP8 training + FP8 rollout | Experimental | Useful for research on training/inference mismatch and throughput, but still has optimizer and checkpointing caveats. | + +## BF16 Training with FP8 Rollout + +This is the main production path in vime. + +You can run FP8 rollout by setting `--hf-checkpoint` to a blockwise-quantized Hugging Face checkpoint. Convert a BF16 checkpoint with: + +```bash +python tools/convert_hf_to_fp8.py \ + --model-dir $BF16_MODEL \ + --save-dir $FP8_MODEL \ + --strategy block --block-size 128 128 \ + --max-workers 4 +``` + +Make sure the converted checkpoint's `config.json` contains the correct `quantization_config`. vime uses that config during weight updates, so the training side can remain BF16 while rollout receives FP8 weights. + +Example: + +```bash +# Megatron training checkpoint remains BF16 / torch_dist. +--ref-load /path/to/model_torch_dist + +# vLLM rollout checkpoint is FP8 Hugging Face. +--hf-checkpoint /path/to/model-fp8-hf +``` + +## FP8 KV Cache for Rollout + +For long-context, multi-turn, or agentic workloads, KV cache capacity is often the bottleneck. Because vLLM arguments are passed through by adding `--vllm-`, you can enable FP8 KV cache directly: + +```bash +--vllm-kv-cache-dtype fp8_e4m3 +``` + +This is a rollout-side setting. It does not change Megatron training precision; it increases effective vLLM KV cache capacity and can allow longer contexts or higher concurrency, subject to the accuracy/performance behavior of your vLLM version and GPU stack. + +## FP8 Training with FP8 Rollout + +vime also supports experimental FP8 training paths. We observed that FP8 training plus FP8 inference can improve inference throughput and reduce training/inference mismatch in some settings. More details are available in [this blog](https://lmsys.org/blog/2025-11-25-fp8-rl/). + +### Quick Start + +1. Convert your Hugging Face model weights to FP8 format using `tools/convert_hf_to_fp8.py`. + +2. Add the FP8 training flags: + +```bash +--fp8-format e4m3 +--fp8-recipe blockwise +# --fp8-param-gather # optional; currently incompatible with CPU Adam +``` + +3. Ensure `NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1` is set. vime sets this to `1` by default for Ray actors. + +4. Start an FP8 training example: + +```bash +# Qwen3-4B FP8 training +bash scripts/low_precision/run-qwen3-4b-fp8.sh + +# Qwen3-30B-A3B FP8 training (2 nodes) +bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh +``` + +### Implementation Notes + +1. If an FP8 recipe is enabled, TransformerEngine layers are built in an FP8 context. +2. During training, weights and activations are quantized online to NVFP8 format, and cuBLAS FP8 GEMM is used for forward and backward GEMMs. +3. During RL weight updates, Megatron dequantizes FP8 weights to BF16, then vime quantizes the BF16 weights to FP8 and sends them to vLLM. +4. Checkpoints saved from the training engine are dequantized back to BF16 and saved as `torch_dist`. + +Only `Linear` and `GroupLinear` layers in TransformerEngine use FP8. `embedding` and `lm_head` remain in their original precision. If `--fp8-param-gather` is not enabled, TransformerEngine weights remain stored in BF16 and are cast to FP8 only during `GEMM` or `GroupGEMM`. + +### Known Caveat + +`--fp8-param-gather` can save memory, but currently requires TransformerEngine `FusedAdam`, which conflicts with the CPU Adam offload path commonly used for large Megatron-LM RL jobs. + +## INT4 QAT Training + +INT4 STE (Straight-Through Estimator) training and INT4 inference can further reduce rollout memory and improve throughput. Treat this path as beta unless you have validated the target model and reward setup. + +### Quick Start + +1. Convert Hugging Face weights to INT4: + +```bash +python tools/convert_hf_to_int4_direct.py \ + --model-dir /path/to/your/original/models \ + --save-dir /path/to/your/save/models +``` + +If you only need INT4 rollout, set `--hf-checkpoint` to the converted INT4 checkpoint. + +2. Enable INT4 fake QAT: + +```json +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" +``` + +`OPEN_TRAINING_INT4_GROUP_SIZE` should usually be: + +- `128` for `moonlight-16B-A3B`, `qwen3-30B-A3B`, and `qwen3-235B-A22B-int4`; +- `32` for `kimi-k2-Thinking-int4`. + +3. Launch an example: + +```bash +# Moonlight-16B-A3B INT4 training +bash scripts/low_precision/run-moonlight-16B-A3B-int4.sh + +# Qwen3-30B-A3B INT4 training +bash scripts/low_precision/run-qwen3-30B-A3B-int4.sh + +# Qwen3-235B-A22B INT4 training (8 nodes) +bash scripts/low_precision/run-qwen3-235B-A22B-int4.sh + +# Kimi-k2-Thinking INT4 training (32 nodes) +bash scripts/low_precision/run-kimi-k2-Thinking-int4.sh +``` + +For multi-node environments, start the Ray service according to your cluster configuration. diff --git a/docs/en/advanced/megatron-config.md b/docs/en/advanced/megatron-config.md index 42a8ad445..d69a49cfb 100644 --- a/docs/en/advanced/megatron-config.md +++ b/docs/en/advanced/megatron-config.md @@ -74,7 +74,6 @@ megatron: ```bash python train.py \ --advantage-estimator ppo \ - --use-critic \ --megatron-config-path megatron_ppo.yaml \ --tensor-model-parallel-size 2 \ --sequence-parallel \ @@ -84,14 +83,13 @@ python train.py \ --expert-tensor-parallel-size 1 \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 8 \ - --critic-num-nodes 1 \ - --critic-num-gpus-per-node 8 \ ... ``` In this setup: -- CLI defines the shared topology and resource layout. +- `--advantage-estimator ppo` enables the critic automatically; there is no separate `--use-critic` CLI flag. +- CLI defines the shared topology and resource layout; in current PPO, critic training resources follow the actor configuration. - YAML defines the role-specific differences, such as `lr`, `load`, `save`, or optimizer / scheduler parameters. ### Overriding Only One Role @@ -114,6 +112,7 @@ In this case the actor keeps the shared CLI arguments unchanged. - **PPO only for now.** `--megatron-config-path` is currently intended for PPO actor / critic role configuration. It is not the recommended interface for GRPO, REINFORCE++, and other critic-free workflows. - **Actor and critic must use the same Megatron parallel topology in current PPO.** In particular, topology-related settings such as `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `context_parallel_size`, `expert_model_parallel_size`, `expert_tensor_parallel_size`, and `sequence_parallel` should not differ between actor and critic. +- **Actor and critic share the same train placement group in current PPO.** The critic node count and GPUs per node are derived from the actor configuration and cannot be used as an independent resource scale. - **Keep topology-related settings on CLI.** The safest current pattern is to keep parallelism and resource arguments in the shared CLI configuration, and only put role-specific differences in YAML, such as `lr`, `load`, `save`, warmup, and optimizer / scheduler settings. If you configure different parallel topologies for actor and critic, the behavior is currently unsupported and may fail during initialization or training. @@ -126,6 +125,6 @@ If you configure different parallel topologies for actor and critic, the behavio Yes. Missing roles automatically inherit the shared CLI arguments, so you do not need to duplicate everything in YAML. -### Q: Can I move `--actor-num-nodes` or `--critic-num-gpus-per-node` into YAML? +### Q: Can I move resource settings into YAML? -No. Resource allocation and placement groups are still controlled by CLI arguments, and the corresponding YAML fields are ignored. \ No newline at end of file +No. Resource allocation and placement groups are still controlled by CLI arguments, and the corresponding YAML fields are ignored. `--actor-num-nodes` / `--actor-num-gpus-per-node` determine the PPO train resource scale; the critic node count and GPUs per node follow the actor configuration and cannot be configured independently. diff --git a/docs/en/advanced/observability.md b/docs/en/advanced/observability.md new file mode 100644 index 000000000..d8c84a62a --- /dev/null +++ b/docs/en/advanced/observability.md @@ -0,0 +1,99 @@ +# Observability + +vime's default observability path is intentionally small: training metrics still go to W&B / TensorBoard; high-frequency vLLM Prometheus metrics are no longer uploaded to W&B; request timings from vLLM response `meta_info` are stored in sample traces and aggregated once per rollout step as compact `perf/...` metrics. + +## W&B / TensorBoard Metrics + +W&B and TensorBoard still receive reward, loss, KL, entropy, eval, and other training metrics. vLLM request timing summaries are logged under `perf/`, for example: + +```text +perf/request/e2e_latency/mean +perf/request/queue_time/median +perf/request/count +perf/request/profiled_count +perf/decode/throughput/mean +perf/prefill/bootstrap_queue_duration/mean +perf/prefill/bootstrap_duration/mean +perf/prefill/alloc_wait_duration/mean +perf/prefill/forward_duration/max +perf/prefill/transfer_speed_gb_s/mean +perf/decode/prealloc_duration/mean +perf/decode/bootstrap_duration/mean +perf/decode/alloc_wait_duration/mean +perf/decode/transfer_duration/max +perf/decode/forward_duration/mean +``` + +These metrics are aggregated once per rollout step, not emitted once per request, so they should not slow W&B like uploading raw Prometheus metrics would. + +Without PD, common `perf/request/...` metrics and available `perf/decode/throughput/...` metrics still exist. Detailed `perf/prefill/...` and `perf/decode/...duration` metrics only appear when vLLM returns the corresponding `pd_*` timing fields. + +## Where Prometheus Data Is Stored + +vime does not store per-second Prometheus data. vLLM / router only expose `/metrics` and `/engine_metrics` HTTP endpoints. Prometheus scrapes those endpoints periodically and stores the time series in Prometheus's own TSDB. + +That means: + +- If Prometheus is not running, serving metrics are only available from the current vLLM process memory and endpoint output, with no historical storage. +- If Prometheus is running, history is stored under Prometheus's `--storage.tsdb.path`. +- vime does not upload these high-frequency metrics to W&B. + +Useful vLLM metrics include: + +```text +vllm:num_queue_reqs +vllm:num_running_reqs +vllm:num_prefill_bootstrap_queue_reqs +vllm:num_prefill_inflight_queue_reqs +vllm:num_decode_prealloc_queue_reqs +vllm:num_decode_transfer_queue_reqs +vllm:kv_transfer_speed_gb_s_bucket +vllm:kv_transfer_latency_ms_bucket +vllm:kv_transfer_total_mb_bucket +``` + +These are useful in Prometheus / Grafana for live queue buildup, transfer speed, latency histograms, failure counters, and other serving-side symptoms. + +## Starting Prometheus + +Prometheus must run while training is running because it can only scrape endpoints that are currently alive. It does not need to run inside the training Python process; run it as a side process in the same machine or job. + +A minimal config is: + +```yaml +global: + scrape_interval: 10s + +scrape_configs: + - job_name: vime-vllm + metrics_path: /engine_metrics + static_configs: + - targets: + - "ROUTER_IP:ROUTER_PORT" +``` + +Replace `ROUTER_IP:ROUTER_PORT` with the router address printed by vime, or with the explicit `--vllm-router-ip` / `--vllm-router-port` values. + +Start Prometheus with its TSDB path on persistent storage: + +```bash +prometheus \ + --config.file=/path/to/prometheus.yml \ + --storage.tsdb.path=/path/to/prometheus-data \ + --storage.tsdb.retention.time=7d \ + --web.listen-address=0.0.0.0:9090 +``` + +The vime image includes the `prometheus` binary, so this command can run directly inside the container. You can also start a side container from the same image as long as it can reach the router address and mounts `/path/to/prometheus-data` on persistent storage. + +If `--storage.tsdb.path` points to container-local disk, the data is lost when the container is removed. If it points to NFS, a persistent volume, or a job output directory, you can restart Prometheus with the same TSDB directory after training and query the historical time range in the Prometheus UI or Grafana. This is time-series replay, not full per-request trace replay; per-sample request timings still come from sample traces / debug rollout data. + +## Trace Viewer + +Debug rollout dumps saved with `--save-debug-rollout-data` include sample traces. The trace viewer reads vLLM timing attrs directly from those traces and uses `pd_*` fields to render synthetic `[P]` / `[D]` lanes. + +```bash +python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt +``` + +The default path does not require separate `ReqTimeStats(...)` logs, Loki, or a compaction tool. diff --git a/docs/en/advanced/on-policy-distillation.md b/docs/en/advanced/on-policy-distillation.md new file mode 100644 index 000000000..c6c20a7e5 --- /dev/null +++ b/docs/en/advanced/on-policy-distillation.md @@ -0,0 +1,128 @@ +# On-Policy Distillation + +On-policy distillation (OPD) trains a student on response tokens sampled from the student's current policy. At every visited prefix, a fixed teacher scores the same next token, providing a dense token-level learning signal along the student's own trajectories. In vime, this signal is a sampled reverse-KL penalty applied to the advantage, so it can be combined with an advantage estimator such as GRPO, PPO, or REINFORCE++. With zero task reward, the same mechanism performs pure distillation. + +## Key Arguments + +| Argument | Description | +|----------|-------------| +| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. | +| `--opd-type` | Type of OPD: `vllm` or `megatron`. Required when `--use-opd` is set. | +| `--opd-kl-coef` | OPD KL penalty coefficient (default: 1.0). Controls the weight of the distillation signal relative to the RL advantage. | +| `--opd-teacher-load` | Path to teacher Megatron checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=vllm`. | +| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. | +| `--opd-teacher-model` | Optional served model name sent to the external VLLM teacher when `--opd-type=vllm`. | + +## How It Works + +Let $\pi_\theta$ denote the student, $\pi_T$ the teacher, and $h_t$ the history before token $a_t$ on a student-generated trajectory. Following [Thinking Machines Lab's definition](https://thinkingmachines.ai/blog/on-policy-distillation/), the per-token reverse KL is + +$$ +D_{\mathrm{KL}}\left(\pi_\theta(\cdot \mid h_t) \| \pi_T(\cdot \mid h_t)\right) += \mathbb{E}_{a_t \sim \pi_\theta(\cdot \mid h_t)}\left[ +\log \pi_\theta(a_t \mid h_t) - \log \pi_T(a_t \mid h_t) +\right]. +$$ + +The order is important: the student is the first argument of the KL, and the expectation is also over the student distribution. The teacher does not generate the training trajectory; it evaluates the token that the student actually sampled. + +vime does not enumerate the full vocabulary to compute this expectation. For each sampled token, it uses the Monte Carlo contribution + +$$ +\hat d_t = \log \pi_\theta(a_t \mid h_t) - \log \pi_T(a_t \mid h_t), +\qquad a_t \sim \pi_\theta(\cdot \mid h_t), +$$ + +and modifies the base advantage as + +$$ +\hat A_t = A_t - \lambda_{\mathrm{opd}}\hat d_t. +$$ + +Here, $A_t$ is the advantage from the configured estimator (or zero for pure distillation), and $\lambda_{\mathrm{opd}}$ is `--opd-kl-coef`. An individual $\hat d_t$ may be negative even though the KL is non-negative in expectation. The policy loss uses $\hat A_t$, which makes the OPD term orthogonal to the choice of GRPO, PPO, REINFORCE++, GSPO, or another supported advantage estimator. + +## Two Teacher Modes + +### VLLM Mode (`--opd-type vllm`) + +The teacher runs on an external VLLM server. Teacher log-probs are obtained during the rollout phase. + +**When to use**: The teacher has a different architecture from the student, or the teacher is too large to load alongside the training model. Because the teacher scores the student's exact token IDs, the teacher and student must still use compatible tokenization and vocabularies. + +**How it works**: +1. An external VLLM server runs the teacher model. +2. During rollout, the custom reward function (`vime.rollout.on_policy_distillation.reward_func`) sends the student's sampled token IDs to the teacher server and obtains the teacher log-probability of those same tokens. +3. The custom post-processing function (`vime.rollout.on_policy_distillation.post_process_rewards`) trims the teacher log-probs to the response span and stores them in `sample.teacher_log_probs`. +4. During training, vime subtracts the sampled log-probability difference, scaled by `--opd-kl-coef`, from the base advantage. + +**Configuration**: +```bash +--use-opd +--opd-type vllm +--opd-kl-coef 1.0 +--custom-rm-path vime.rollout.on_policy_distillation.reward_func +--custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards +--rm-url http://:/inference/v1/generate +``` + +### Megatron Mode (`--opd-type megatron`) + +The teacher model is loaded directly into Megatron via `--opd-teacher-load`. Teacher log-probs are computed during the training forward pass. + +**When to use**: The teacher has the same architecture as the student/reference model and fits in GPU memory. + +**How it works**: +1. The teacher model is loaded as an additional Megatron model during initialization. +2. During the training forward pass, the teacher model computes log-probs for each sample. +3. The KL penalty is computed inline and applied to advantages. + +**Configuration**: +```bash +--use-opd +--opd-type megatron +--opd-kl-coef 1.0 +--opd-teacher-load /path/to/teacher_torch_dist +``` + +> **Note**: The teacher checkpoint must be in Megatron format (`torch_dist` or `torch`). You can convert from HuggingFace format using `tools/convert_hf_to_torch_dist.py`. + +## Running the Examples + +Complete example scripts are provided in `examples/on_policy_distillation/`: + +### VLLM Teacher + +```bash +# 1. Download models and data +hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B +hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k + +# 2. Convert student model +cd /root/vime +source scripts/models/qwen3-8B.sh +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-8B \ + --save /root/Qwen3-8B_torch_dist + +# 3. Run +bash examples/on_policy_distillation/run-qwen3-8B-opd.sh +``` + +### Megatron Teacher + +```bash +# 1. Convert both student and teacher models to Megatron format +# 2. Run +bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh +``` + +## Preliminary Results + +Using Qwen3-8B-Base model SFT-ed on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, on-policy distillation with a Qwen3-32B teacher on the remaining data yields: + +| | Pass@1 | +|-----------------------------------------------|--------| +| Qwen3-8B-Base + SFT | 76% | +| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% | diff --git a/docs/en/advanced/pd-disaggregation.md b/docs/en/advanced/pd-disaggregation.md index 6335c756e..b9e8bd921 100644 --- a/docs/en/advanced/pd-disaggregation.md +++ b/docs/en/advanced/pd-disaggregation.md @@ -10,7 +10,7 @@ Use PD Disaggregation when: - decode dominates rollout time; - prefix-cache locality matters for multi-turn sessions; - prefill and decode need different TP, memory, or runtime settings; -- you want an vLLM serving topology that is closer to production serving rather than a single uniform inference group. +- you want a vLLM serving topology that is closer to production serving rather than a single uniform inference group. For short single-turn tasks, the default regular vLLM engine layout is usually simpler. @@ -30,7 +30,7 @@ This is the lightweight path used by simple scripts. It is convenient when you o ### Advanced Path: `--vllm-config` -For production rollout topologies, use [vLLM Config](vllm-config.md). It lets you configure prefill and decode groups independently, and can also express EPD-style layouts, heterogeneous engine groups, multi-model serving, and per-group vLLM overrides. +For production rollout topologies, use [vLLM Config](vllm-config.md). It lets you configure prefill and decode groups independently, and can also express EPD-style layouts, heterogeneous server groups, multi-model serving, and per-group vLLM overrides. Example: @@ -43,12 +43,13 @@ vllm: num_gpus: 4 num_gpus_per_engine: 2 overrides: - chunked_prefill_size: 8192 + enable_chunked_prefill: true + max_num_batched_tokens: 8192 - worker_type: decode num_gpus: 12 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.88 + gpu_memory_utilization: 0.88 ``` Launch with: diff --git a/docs/en/advanced/reproducibility.md b/docs/en/advanced/reproducibility.md index 276b8c343..7e3eda77c 100644 --- a/docs/en/advanced/reproducibility.md +++ b/docs/en/advanced/reproducibility.md @@ -49,3 +49,29 @@ bash scripts/run-qwen2.5-0.5B-reproducibility.sh ``` For screen shots of the wandb, please refer to [pull#370](https://github.com/THUDM/slime/pull/370). + +## Train/rollout log-prob alignment (GLM-5) + +Beyond single-side bitwise reproduction, vime can align the training log-probs with the rollout (inference) log-probs. This is currently supported only for the **GLM-5 structure** (MLA + DSA sparse attention), and requires the deterministic VLLM / batch-invariant DeepGEMM / DeepEP build. Vime installs the required Megatron-side alignment hooks at runtime; no extra Megatron patch is required. + +Supported in this path: + +- DSA sparse attention (`flashmla_sparse` prefill/decode), including deterministic NSA RadixCache/prefix cache; +- DeepGEMM batch-invariant block-FP8 forward for dense and grouped-MoE layers (with BF16 backward); +- fp32 MoE router (the LM head stays bf16 on both train and rollout — matching precision, not fp32, is what aligns); +- VLLM DeepEP low-latency rollout plus Megatron DeepEP normal training. A + compact second normal dispatch preserves every top-k route, and the token + owner performs the weighted reduction in slot order and FP32. Ordinary + Megatron all-to-all is not an alignment backend for this path; +- bf16 or FP8-E4M3 KV cache. For `flashmla_sparse`, VLLM stores packed FP8 + cache entries and gathers/dequantizes only the selected pages before its BF16 + sparse kernel. The maintained gate defaults to FP8-E4M3 and does not use + rollout routing replay (R3), so all main-model parameters, including the + router and experts, execute backward. The auxiliary DSA indexer remains + frozen through `--freeze-indexer`. + +The regression gate is `tests/test_glm52_6layer_deterministic_e2e.py` (6-layer GLM-5.2, single-node EP8): it runs a real Megatron→VLLM online-weight-update rollout, trains all main-model parameters, and asserts `train_rollout_logprob_abs_diff < 1e-6` (the established DeepEP alignment reference is in the `x e-7` range). + +An additional short EP8 gate, `tests/test_glm52_layerwise_zero_e2e.py`, records +the visible output of decoder layers 0–5 on both sides and requires every +matched hidden-state element to have an absolute difference of exactly zero. diff --git a/docs/en/advanced/speculative-decoding.md b/docs/en/advanced/speculative-decoding.md index 1d1e34636..6cc4188ce 100644 --- a/docs/en/advanced/speculative-decoding.md +++ b/docs/en/advanced/speculative-decoding.md @@ -9,14 +9,14 @@ which vime forwards via `--vllm-speculative-config`. For models with MTP layers (e.g., GLM-4.7, DeepSeek-V3/R1), pass: ```bash ---vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' +--vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' ``` To use a separately trained draft model, set `model` (and optionally `draft_tensor_parallel_size`) in the same JSON: ```bash ---vllm-speculative-config '{"method":"eagle","num_speculative_tokens":3,"model":"/your/draft/model/path"}' +--vllm-speculative-config '{"method":"eagle","num_speculative_tokens":4,"model":"/your/draft/model/path"}' ``` To train a draft model from scratch, see [TorchSpec](https://github.com/lightseekorg/TorchSpec) diff --git a/docs/en/advanced/vllm-config.md b/docs/en/advanced/vllm-config.md index 21b7975c6..aae8e2654 100644 --- a/docs/en/advanced/vllm-config.md +++ b/docs/en/advanced/vllm-config.md @@ -31,11 +31,11 @@ vllm: - name: # Required. Unique identifier for this model. model_path: # Optional. HF checkpoint path. Defaults to --hf-checkpoint. update_weights: # Optional. Whether to sync weights from training. Auto-inferred. - num_gpus_per_engine: # Optional. Default TP size for all groups in this model. + num_gpus_per_engine: # Optional. Default worker GPU count per engine. server_groups: # Required. List of server group configurations. - worker_type: # Required. One of: regular, prefill, decode, placeholder. num_gpus: # Required. Total GPUs allocated to this group. - num_gpus_per_engine: # Optional. TP size override for this group. + num_gpus_per_engine: # Optional. Worker GPU count override for this group. overrides: # Optional. vLLM EngineArgs field overrides. ``` @@ -48,7 +48,7 @@ vllm: | `name` | `str` | **Required** | Unique name for this model (e.g., `"actor"`, `"ref"`, `"reward"`). Used as the key in `args.vllm_model_routers`. | | `model_path` | `str` | `args.hf_checkpoint` | HuggingFace checkpoint path. All server groups within a model must use the same model path. | | `update_weights` | `bool` | Auto | Whether this model receives weight updates from training. When not set, automatically inferred: `true` if `model_path` matches `--hf-checkpoint`, `false` otherwise. | -| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | Default TP size for server groups in this model. Individual groups can override. | +| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | Default total worker GPU count per engine. Individual groups can override. | | `server_groups` | `list` | **Required** | List of `ServerGroupConfig` entries defining the engine topology. (`engine_groups` is accepted as a backward-compatible alias.) | #### Server Group Fields @@ -57,7 +57,7 @@ vllm: |-------|------|---------|-------------| | `worker_type` | `str` | **Required** | Engine type: `regular` (standard), `prefill` (PD prefill worker), `decode` (PD decode worker), or `placeholder` (reserve GPU slots without launching engines). | | `num_gpus` | `int` | **Required** | Total number of GPUs for this group. Must be > 0. | -| `num_gpus_per_engine` | `int` | Model's `num_gpus_per_engine` | TP size override. Number of GPUs per engine instance. | +| `num_gpus_per_engine` | `int` | Model's `num_gpus_per_engine` | Total worker GPU count per engine instance. This equals TP only when DP and PP are both 1. | | `overrides` | `dict` | `{}` | vLLM `EngineArgs` field overrides. Applied on top of `--vllm-*` CLI args with highest priority. | ### Worker Types @@ -107,10 +107,10 @@ vllm: server_groups: - worker_type: prefill num_gpus: 4 - num_gpus_per_engine: 2 # 2 prefill engines, TP=2 + num_gpus_per_engine: 2 # 2 prefill engines, TP=2 with default DP/PP - worker_type: decode num_gpus: 12 - num_gpus_per_engine: 4 # 3 decode engines, TP=4 + num_gpus_per_engine: 4 # 3 decode engines, TP=4 with default DP/PP ``` ```bash @@ -174,18 +174,25 @@ from vime.rollout.vllm_rollout import get_model_url from vime.utils.http_utils import post async def my_generate(args, sample, sampling_params): - # Route to the actor model (default) - actor_url = get_model_url(args, "actor", "/generate") - output = await post(actor_url, {"text": sample.prompt, "sampling_params": sampling_params}) - + # Route to the actor model (default endpoint is /inference/v1/generate) + actor_url = get_model_url(args, "actor") + output = await post(actor_url, { + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, + }) + # output["choices"][0] carries token_ids, logprobs.content[i].logprob, and finish_reason + # Route to the reference model - ref_url = get_model_url(args, "ref", "/generate") - ref_output = await post(ref_url, {"text": sample.prompt, "sampling_params": sampling_params}) - + ref_url = get_model_url(args, "ref") + ref_output = await post(ref_url, { + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, + }) + # Route to the reward model (e.g., OpenAI-compatible API) reward_url = get_model_url(args, "reward", "/v1/chat/completions") reward_output = await post(reward_url, {...}) - + ... ``` @@ -232,9 +239,9 @@ vllm: num_gpus: 2 # reserve 2 GPUs (no engines created) ``` -### 6. Per-Group ServerArgs Overrides +### 6. Per-Group EngineArgs Overrides -Use `overrides` to apply vLLM `ServerArgs` fields to specific server groups without affecting others: +Use `overrides` to apply vLLM `EngineArgs` fields to specific server groups without affecting others: ```yaml vllm: @@ -244,10 +251,12 @@ vllm: num_gpus: 8 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.85 - context_length: 32768 - chunked_prefill_size: 4096 - enable_torch_compile: true + gpu_memory_utilization: 0.85 + max_model_len: 32768 + enable_chunked_prefill: true + max_num_batched_tokens: 4096 + compilation_config: + mode: 3 ``` Overrides take **highest priority**, overriding both the base `--vllm-*` CLI args and model-level defaults. This is especially useful for: @@ -257,7 +266,7 @@ Overrides take **highest priority**, overriding both the base `--vllm-*` CLI arg ### 7. Standalone vLLM Launcher -While `--vllm-config` is designed for vime's training pipeline, it also works as a powerful launcher for pure inference scenarios using the `--rollout-external` pattern or by configuring vime to focus solely on serving. +While `--vllm-config` is designed for vime's training pipeline, it also works as a powerful launcher for pure inference scenarios using external engine addresses or by configuring vime to focus solely on serving. **Using external engines with a pre-launched topology:** @@ -265,17 +274,24 @@ For complex production deployments, you may want to pre-launch vLLM engines inde ```bash # Step 1: Launch vLLM engines externally -vllm serve /path/to/model --port 10090 ... -vllm serve /path/to/model --port 10091 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10090 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10091 ... # Step 2: Connect vime to external engines python train.py \ - --rollout-external \ --rollout-external-engine-addrs host1:10090 host2:10091 \ ... ``` -> **Note:** `--vllm-config` and `--rollout-external` are mutually exclusive. Use `--vllm-config` when you want vime to manage the full engine lifecycle; use `--rollout-external` when engines are pre-deployed. +vime queries each external engine's `/server_info` endpoint to infer +`rollout_num_gpus`, per-engine GPU counts, vLLM parallel sizes, and +prefill/decode worker types. If no `--vllm-router-ip/--vllm-router-port` +is provided, vime launches its own router and registers the external engines +to it. + +> **Note:** `--vllm-config` and `--rollout-external-engine-addrs` are mutually exclusive. Use `--vllm-config` when you want vime to manage the full engine lifecycle; use `--rollout-external-engine-addrs` when engines are pre-deployed. + +For external-engine selection, update from disk, and delta disk transport, see [External Rollout Engines Roadmap](external-rollout-engines.md). --- @@ -332,7 +348,7 @@ When the config is loaded, vime applies the following resolution cascade: | Flag | Conflict Reason | |------|----------------| | `--prefill-num-servers` | PD disaggregation is configured via `server_groups` in the YAML | -| `--rollout-external` | External engines have their own topology; config manages the lifecycle internally | +| `--rollout-external-engine-addrs` | External engines have their own topology; config manages the lifecycle internally | --- @@ -351,12 +367,13 @@ vllm: num_gpus: 4 num_gpus_per_engine: 2 overrides: - chunked_prefill_size: 8192 + enable_chunked_prefill: true + max_num_batched_tokens: 8192 - worker_type: decode num_gpus: 12 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.88 + gpu_memory_utilization: 0.88 - name: ref model_path: /data/models/Qwen3-32B @@ -392,33 +409,42 @@ python train.py \ **Custom rollout function (`my_agent/rollout.py`):** ```python +from transformers import AutoTokenizer + from vime.rollout.vllm_rollout import get_model_url from vime.utils.http_utils import post async def generate_with_models(args, sample, sampling_params): """Generate using actor, score with reward model, compare with reference.""" - # Generate from actor - actor_url = get_model_url(args, "actor", "/generate") + # Generate from actor (default endpoint is /inference/v1/generate) + actor_url = get_model_url(args, "actor") actor_output = await post(actor_url, { - "text": sample.prompt, - "sampling_params": sampling_params, - "return_logprob": True, + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, }) - - # Get reference logprobs for KL penalty - ref_url = get_model_url(args, "ref", "/generate") + response_ids = actor_output["choices"][0]["token_ids"] + + # Get reference logprobs over the prompt+response. max_tokens=1 + prompt_logprobs scores + # the submitted token_ids; read them from the top-level "prompt_logprobs" field. + ref_url = get_model_url(args, "ref") ref_output = await post(ref_url, { - "text": sample.prompt + actor_output["text"], - "sampling_params": {"max_new_tokens": 0, "temperature": 0}, - "return_logprob": True, + "token_ids": sample.tokens + response_ids, + "sampling_params": {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 1}, }) - - # Score with reward model + + # Score the actor response with the reward model (OpenAI-compatible) + tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + response_text = tokenizer.decode(response_ids, skip_special_tokens=True) + prompt_messages = ( + sample.prompt + if isinstance(sample.prompt, list) + else [{"role": "user", "content": sample.prompt}] + ) reward_url = get_model_url(args, "reward", "/v1/chat/completions") reward_output = await post(reward_url, { "model": "reward", - "messages": [{"role": "user", "content": sample.prompt + actor_output["text"]}], + "messages": [*prompt_messages, {"role": "assistant", "content": response_text}], }) # ... process outputs and return Sample @@ -446,7 +472,7 @@ Use `get_model_url(args, "model_name", "/endpoint")` from `vime.rollout.vllm_rol ### Q: Can I use `--vllm-config` without training (inference only)? -While `--vllm-config` is designed for vime's training loop, you can effectively use it for inference-only scenarios by configuring a rollout-only run. For fully standalone vLLM serving, consider using vLLM's native `vllm serve` directly or the `--rollout-external` mode for connecting to pre-deployed engines. +While `--vllm-config` is designed for vime's training loop, you can effectively use it for inference-only scenarios by configuring a rollout-only run. For fully standalone vLLM serving, use the public `vllm serve` command directly or `--rollout-external-engine-addrs` to connect to pre-deployed engines. ### Q: What is the relationship between `--vllm-config` and `--prefill-num-servers`? diff --git a/docs/en/developer_guide/ci.md b/docs/en/developer_guide/ci.md index 60c3609a0..0d40b3083 100644 --- a/docs/en/developer_guide/ci.md +++ b/docs/en/developer_guide/ci.md @@ -1,122 +1,48 @@ # CI (Continuous Integration) -vime uses GitHub Actions for CI. Tests are triggered by **PR labels** — adding a specific label to a PR will run the corresponding test suite. +Vime uses Buildkite for continuous integration. The committed pipeline is +`.buildkite/pipeline.yml`. -## How It Works +## Always-on checks -The workflow is defined in `.github/workflows/pr-test.yml` (auto-generated from `pr-test.yml.j2`). Each CI job: +Every pull request runs these CPU steps: -1. Runs on a self-hosted GPU runner via `docker run`; most tests use `inferactinc/public:vime-latest`, while image validation uses `inferactinc/public:vime-test-latest`. -2. Installs vime with `pip install -e . --no-deps`. -3. Acquires the required GPUs via `tests/ci/gpu_lock_exec.py --count `. -4. Executes the test file: `python .py` or `python tests/.py`, depending on whether the test lives under `tests/` or a subdirectory such as `tests/plugin_contracts/`. +| Step | Coverage | +|---|---| +| `pre-commit` | formatting, lint, and repository policy | +| `plugin-contracts` | customization contracts and CPU tests | +| `agent-adapter` | agent adapter behavior | +| `upstream-sync-cpu` | CPU tests synchronized from upstream | +| `utils` | `tests/utils` | -Each test file follows a standard pattern: a `prepare()` function downloads models/datasets, and an `execute()` function builds CLI arguments and calls `U.execute_train(...)`. +The authoritative commands and queue configuration are in +`.buildkite/pipeline.yml`. -## CI Labels +## GPU suites -Add a label to your PR to trigger the corresponding test suite: +After the CPU steps pass, the Buildkite build exposes a block step named +`Run GPU test suites?`. Select one or more suites: -| Label | Job | Description | -|---|---|---| -| `run-ci-short` | `e2e-test-short` | Lightweight smoke tests with Qwen2.5-0.5B (4 GPUs). Fast feedback loop. | -| `run-ci-megatron` | `e2e-test-megatron` | Core Megatron training tests covering dense, MoE, PPO, MTP, etc. | -| `run-ci-precision` | `e2e-test-precision` | Numerical precision validation (parallel check). | -| `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint save/load correctness (sync and async-save). | -| `run-ci-image` | `e2e-test-image` | Full test suite run on `inferactinc/public:vime-test-latest` image (for image validation). | -| `run-ci-changed` | `e2e-test-changed` | **Dynamically** detects new/modified test files in the PR and runs only those. | +- `short` +- `vllm-config` +- `megatron` +- `vime-customized` +- `precision` +- `ckpt` -All labels also run when triggered via `workflow_dispatch` (manual run from the Actions tab). +`.buildkite/gpu_suites.py` expands each selected suite into one Buildkite job +per test. Set `VIME_CI_IMAGE` to an immutable candidate digest when validating +Dockerfile or vLLM patch changes. Jobs otherwise use `vllm/vime:latest`, which +must not be updated before the change merges. -## Key Labels Explained +## Registering tests -### `run-ci-changed` — Run Only New or Modified Tests +- Add always-on CPU tests to the appropriate command in + `.buildkite/pipeline.yml`. +- Add GPU tests to a suite in `.buildkite/gpu_suites.py` and update the suite + count shown by `.buildkite/pipeline.yml`. +- Keep `.buildkite/README.md` synchronized with pipeline behavior. -This is the most useful label for development. When you add a new test file or modify an existing one, just add `run-ci-changed` to your PR and CI will: - -1. **Detect** which `tests/test_*.py` or `tests/plugin_contracts/test_*.py` files are added or modified relative to `origin/main` (via `git diff --diff-filter=AM`). -2. **Extract** the `NUM_GPUS` value from each detected test file automatically. -3. **Build** a dynamic GitHub Actions matrix and run each test in parallel. - -This means you don't need to manually register your new test in the workflow — just make sure your test file has a top-level `NUM_GPUS = ` constant and `run-ci-changed` will pick it up. - -**Example**: If your PR adds `tests/test_mimo_7B_mtp_only_grad.py` with `NUM_GPUS = 8`, adding the `run-ci-changed` label will automatically run that test on 8 GPUs. - -### `run-ci-image` — Full Suite on Test Image - -This runs **all** registered tests on the `inferactinc/public:vime-test-latest` Docker image. Use this label to: - -- Validate a newly built Docker image before release. -- Run the entire test suite for a comprehensive pre-merge check. - -Since this includes every test, it consumes significant GPU time — use it sparingly and prefer more targeted labels for routine development. - -### `run-ci-megatron` — Core Megatron Tests - -This is the primary label for validating Megatron-backend changes. It covers: - -- Dense models: GLM4-9B, Qwen3-4B (PPO) -- MoE models: Qwen3-30B-A3B (with DeepEP), Qwen3.6-35B-A3B PD + Mooncake, Moonlight-16B-A3B -- Specialized: MiMo-7B MTP, Qwen2.5-0.5B debug rollout-then-train - -All tests use 8 GPUs. If you are modifying Megatron training logic, loss computation, or checkpoint conversion, this is the label to use. - -## Writing a New Test - -1. Create `tests/test_.py` following the standard pattern: - -```python -import os -import vime.utils.external_utils.command_utils as U - -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" -NUM_GPUS = 4 # This constant is used by run-ci-changed - -def prepare(): - U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - # Download datasets as needed ... - -def execute(): - # Build argument strings and call U.execute_train(...) - ... - -if __name__ == "__main__": - prepare() - for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): - os.environ.pop(proxy_var, None) - execute() -``` - -2. **For quick validation**: Just push your test file and add `run-ci-changed` to the PR. It will be auto-detected. - -3. **To register in a permanent label group**: Edit `.github/workflows/pr-test.yml.j2`, add an entry to the desired job's `tests` list, then regenerate: - -```bash -cd .github/workflows && python generate_github_workflows.py -``` - -Remember to commit both the `.j2` and the generated `.yml` file. - -## Workflow Generation - -The workflow file `pr-test.yml` is auto-generated from the Jinja2 template `pr-test.yml.j2`. **Do not edit `pr-test.yml` directly.** To make changes: - -1. Edit `.github/workflows/pr-test.yml.j2`. -2. Run `python .github/workflows/generate_github_workflows.py`. -3. Commit both files. - -## Customization Contract Tests - -For CPU-only contract tests that validate hooks loaded from function paths, run: - -```bash -python -m pytest \ - tests/plugin_contracts/test_plugin_rollout_contracts.py \ - tests/plugin_contracts/test_plugin_generate_contracts.py \ - tests/plugin_contracts/test_plugin_path_loading_contracts.py \ - tests/plugin_contracts/test_plugin_runtime_hook_contracts.py -``` - -These files also support direct execution as `python tests/plugin_contracts/.py`. They declare `NUM_GPUS = 0`, so `run-ci-changed` can pick them up without treating them as GPU-heavy end-to-end tests. +Run the exact command locally before triggering its remote Buildkite job. For +GPU failures, reproduce on an H200 node with the same image and environment, +then rerun the remote suite only after the local test passes. diff --git a/docs/en/developer_guide/debug.md b/docs/en/developer_guide/debug.md index 212affd16..d60f4a6fd 100644 --- a/docs/en/developer_guide/debug.md +++ b/docs/en/developer_guide/debug.md @@ -50,6 +50,54 @@ Specifically, vime currently provides the following parameters for separate debu When enabled, data will be loaded from `args.load_debug_rollout_data.format(rollout_id=rollout_id)`, and vLLM will not be initialized (automatically setting `debug_train_only=True`). This method allows you to fix the input for the training part to tune it, for example, by switching between different parallelization strategies. +5. `--save-debug-train-data /your/saved/debug/train_{rollout_id}.pt` + + Saves one train-side file per rollout. Only the last Pipeline Parallel stage and Tensor Parallel rank 0 participate. They restore response-token fields such as `log_probs`, `ref_log_probs`, `values`, `advantages`, `returns`, `kl`, and `entropy` across Context Parallel ranks. Context Parallel rank 0 moves each restored tensor to CPU immediately, so complete tensors do not accumulate on the GPU, and then gathers the distinct Data Parallel shards to one writer. + + The version-2 payload mirrors the rollout debug dump: a top-level `samples` list holds one dict per training sample (`sample_index`, `data_parallel_rank`, and its per-sample fields such as `tokens`, `log_probs`, `advantages`), sorted by `sample_index` so it lines up one-to-one with the rollout dump's `samples` (join on `sample_index` ↔ the rollout side's `index`). A parallel `dp_shards` key preserves the DP/micro-batch layout — each entry records `rank`, `data_parallel_rank`, that shard's `sample_indices`, and the DP-local schedule (`micro_batch_indices`, `num_microbatches`, `global_batch_sizes`) — without duplicating any per-sample tensor. Whole-batch fields such as `raw_reward` are stored once at the top level. If any sample lacks a `sample_index` (custom rollouts that build fresh `Sample` objects leave it `None`), the samples stay in DP-gather order and a warning is logged. With or without CP, response-token fields use the same full-response format. In configs that skip the separate actor log-prob recompute (`can_reuse_log_probs_in_loss` or `--use-rollout-logprobs`), the actor `log_probs` are snapshotted from the training forward itself (keyed by rollout position, at no extra forward), so the dump still carries them. + +## INT4 / Compressed-Tensors Quantization Checkpoint Issues + +When using INT4-quantized models (e.g., `compressed-tensors` with `W4A16`), the checkpoint's `config.json` contains a `quantization_config.ignore` list that specifies which parameters should **not** be quantized. During online weight updates (Megatron → vLLM), vime also reads this ignore list to decide which parameters to INT4-quantize. An incorrect ignore list can cause silent errors: + +1. **MoE router weights (`mlp.gate.weight`) become all zeros** + + The MoE router weight (`mlp.gate.weight`, shape `[num_experts, hidden_size]`) is a plain 2D weight tensor, but it is **not** a Linear layer weight. If it is not in the ignore list, the online quantizer will INT4-quantize it into `weight_packed`, `weight_scale`, `weight_zero_point`, etc. However, vLLM does not expect quantized names for the router, so these parameters are silently skipped during `load_weights`, resulting in all-zero gate weights. + + **Fix**: Ensure `config.json` contains `"re:.*mlp\\.gate\\..*"` in the ignore list. + +2. **Other non-Linear 2D weights** + + Similar issues can occur with any 2D `.weight` tensor that is not a true Linear layer, such as `model.embed_tokens.weight`. Always verify the ignore list covers all non-Linear weights. + + **Recommended ignore patterns** (for GLM-style MoE models): + ```json + "ignore": [ + "lm_head", + "model.embed_tokens.weight", + "re:.*self_attn.*", + "re:.*mlp\\.shared_experts.*", + "re:.*mlp\\.gate_up_proj.*", + "re:.*mlp\\.gate_proj.*", + "re:.*mlp\\.up_proj.*", + "re:.*mlp\\.down_proj.*", + "re:.*eh_proj.*", + "re:.*mlp\\.gate\\..*" + ] + ``` + +3. **Missing safetensors shards** + + Conversion tools may occasionally produce an incomplete checkpoint (e.g., a missing `model-00010-of-00093.safetensors`). After conversion, always verify: + - The number of `.safetensors` files matches the expected count. + - The `model.safetensors.index.json` contains entries for every layer. + - Spot-check that critical layers (e.g., the first MoE layer) have the expected number of keys. + +4. **How to diagnose** + + - Use `--check-weight-update-equal` to verify that weights after a Megatron → vLLM sync match the expected values. If a parameter shows all zeros on the vLLM side, it was likely incorrectly quantized or missing from the checkpoint. + - Use `--debug-rollout-only` with a small number of GPUs to quickly test whether vLLM can generate coherent text from the quantized checkpoint alone. + ## Debug vllm illegal memory access (IMA) When running large scale RL, we will occationally meet the IMA in vLLM, there are some debug suggestions based on our experience: diff --git a/docs/en/developer_guide/profiling.md b/docs/en/developer_guide/profiling.md index 4db02cf8b..29168ea90 100644 --- a/docs/en/developer_guide/profiling.md +++ b/docs/en/developer_guide/profiling.md @@ -47,18 +47,18 @@ Common JSON fields: | `max_iterations` | Worker auto-stops and flushes after more than N steps (condition is `> N`) | | `ignore_frontend` | Recommended `true`: profile workers only, lower frontend overhead | -**Avoid RPC timeout on `stop_profile`:** vLLM APIServer talks to EngineCore/workers over internal RPC. Manually calling `stop_profile` to flush traces can take minutes, while the default `VLLM_RPC_TIMEOUT` is only **10 seconds** (10000 ms), which can interrupt flush or leave traces incomplete. For profiling, set **30 minutes** (1800000 ms). +The `/stop_profile` request waits for EngineCore and all workers to finish +flushing their traces. Large profiles can therefore take several minutes; do +not interrupt the request while the trace files are being written. -Set this variable **before starting train and launching vLLM**, in the Ray worker environment (a local shell `export` may not reach the Ray job). Pass it via `runtime-env-json` on `ray job submit`, for example: +Pass the regular worker environment through `runtime-env-json` on +`ray job submit`, for example: ```bash -export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" } }" @@ -124,8 +124,8 @@ python tools/profile_rollout.py \ While `sleep_rollout` is waiting: 1. `profile_rollout.py --action start` -2. Send a few completion requests to the router or **directly to a worker** (2–4 is enough; traces get large) -3. (Optional) `profile_rollout.py --action stop`; or wait for `max_iterations` to auto-flush +2. Send a few completion requests to the router or **directly to a worker** (2-4 is usually enough; traces get large) +3. If relying on auto-flush, remember that `max_iterations` stops after `> N` steps. For example, `max_iterations=3` needs 4 requests; otherwise call `profile_rollout.py --action stop` manually. 4. Inspect traces under `torch_profiler_dir` Example request (`model` is the HF checkpoint path): @@ -162,9 +162,9 @@ python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-rank | Symptom | Fix | |------|------| | `POST /start_profile` 404 | Pass `--vllm-profiler-config` as JSON; restart the job | -| Start OK but empty output dir | Confirm curl hits a worker and returns 200; increase `max_iterations` or send more requests | +| Start OK but empty output dir | Confirm curl hits a worker and returns 200; if `max_iterations=3`, send 4 requests or call `stop_profile` manually | | Router 503 | Confirm the current job's router port; connect directly to a worker | -| Slow or timed-out stop | Increase `VLLM_RPC_TIMEOUT`; reduce request count | +| Slow stop | Wait for trace flushing to finish; reduce request count | ## 8. Full Runnable Example @@ -212,13 +212,10 @@ launch_train_for_profiling() { source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" - export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" } }" @@ -290,16 +287,15 @@ run_profiling_session() { echo "=== 1/3 start_profile (all workers via router) ===" python tools/profile_rollout.py --router-url "${router_url}" --action start - echo "=== 2/3 send completions (direct to worker; 3 requests) ===" - for i in 1 2 3; do - curl -sS -X POST "${worker_url}/v1/completions" \ + echo "=== 2/3 send completions (direct to worker; 4 requests so max_iterations=3 can auto-flush) ===" + for i in 1 2 3 4; do + response="$(curl -sS -X POST "${worker_url}/v1/completions" \ -H "Content-Type: application/json" \ - -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ - | head -c 400 - echo + -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}")" + printf '%s\n' "${response:0:400}" done - echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" + echo "=== 3/3 list trace files (max_iterations=3 auto-stop uses > N; add --action stop if needed) ===" sleep 2 find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" diff --git a/docs/en/developer_guide/trace.md b/docs/en/developer_guide/trace.md index 31f3c8163..29f3ee8f9 100644 --- a/docs/en/developer_guide/trace.md +++ b/docs/en/developer_guide/trace.md @@ -41,7 +41,7 @@ By default it also starts a local static server so you can open the generated HT ## Instrument custom code -For custom rollout or reward code — including custom agent steps, tool calls, sandbox execution, and verifier calls in agentic workflows — reuse helpers from `vime.utils.trace_utils`: +For custom rollout or reward code — including custom agent steps, tool calls, sandbox execution, and verifier calls in agentic workflows — reuse helpers from `vime.observability.trace_utils`: - `trace_span(target, name, attrs=...)`: record a duration span. - `trace_event(target, name, attrs=...)`: record an instant event. @@ -57,7 +57,7 @@ Use `trace_function(...)` when the whole function should be represented as one s The decorator is what vime uses for the main rollout pipeline. For example, `generate_and_rm(...)` is traced per sample and `generate_and_rm_group(...)` is traced per sample group: ```python -from vime.utils.trace_utils import trace_function +from vime.observability.trace_utils import trace_function @trace_function("generate_and_rm", target="sample") @@ -104,7 +104,7 @@ If you need to add attrs after part of the function has executed, use an inner ` If you want to record vLLM generation metadata in a consistent way, reuse `build_vllm_meta_trace_attrs`: ```python -from vime.utils.trace_utils import build_vllm_meta_trace_attrs, trace_span +from vime.observability.trace_utils import build_vllm_meta_trace_attrs, trace_span with trace_span(sample, "vllm_generate") as span: output = await post(url, payload) @@ -116,4 +116,3 @@ with trace_span(sample, "vllm_generate") as span: - Save a small number of rollouts first; the viewer is easiest to read when each dump contains a manageable number of samples. - The viewer is built from the saved `.pt` dump, so traces can be inspected offline on another machine. - For GPU/kernel-level vLLM profiling traces, see [Profiling](./profiling.md). - diff --git a/docs/en/examples/deepseek-r1.md b/docs/en/examples/deepseek-r1.md new file mode 100644 index 000000000..5b327be1a --- /dev/null +++ b/docs/en/examples/deepseek-r1.md @@ -0,0 +1,206 @@ +# DeepSeek R1 with 128xH100 + +This is an example of doing DeepSeek R1 RL training using 128xH100 GPUs. + +We will use bf16 for training, and an fp8 format with 128x128 blockwise quantization for inference. The maximum response length is 32k, and dynamic sampling will be used to filter data during training. + +Regarding parallelism, for vLLM we will enable expert parallelism (`--vllm-enable-expert-parallel`) and data parallelism (`--vllm-data-parallel-size 8`). DeepEP is disabled by default. For the Megatron part, we will use TP8, PP4, EP32, and CP4. + +⚠️ To save GPU memory, we will use CPU Adam. Each node (8xH100) will occupy 1.4\~1.5TB of host memory. If a single machine's host memory is insufficient, this can be resolved by adding more GPUs to expand the parallelism. + +## Environment Setup + +For instructions on setting up the environment and downloading data, please refer to [Example: Qwen3-4B](qwen3-4B.md). + +To prepare the DeepSeek R1 checkpoint, first you will need to download DeepSeek-R1 to a directory accessible by all machines (hereinafter referred to as `$BASE_DIR`): + +```bash +hf download deepseek-ai/DeepSeek-R1 --local-dir $BASE_DIR/DeepSeek-R1 +``` + +The Hugging Face checkpoint for DeepSeek-R1 is in a block-quantized fp8 format. To convert it into a torch_dist format that Megatron can load, you first need to convert it to a bf16 Hugging Face checkpoint: + +```bash +cd vime/ +python tools/fp8_cast_bf16.py --input-fp8-hf-path $BASE_DIR/DeepSeek-R1 --output-bf16-hf-path $BASE_DIR/DeepSeek-R1-bf16/ +``` + +Next, we need to convert the bf16 version of DeepSeek-R1 into the torch_dist format. Specifically, execute the following on 4 separate nodes: + +```bash +cd vime/ +source scripts/models/deepseek-v3.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=4 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 8 \ + --expert-tensor-parallel-size 1 \ + --expert-model-parallel-size 4 \ + --decoder-first-pipeline-num-layers 7 \ + --decoder-last-pipeline-num-layers 6 \ + --hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ \ + --save $BASE_DIR/DeepSeek-R1_torch_dist/ +``` + +Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` indicates the node's index, both configured similarly to a multi-node `torchrun` setup. + +## Executing the Training + +On node0, run: + +```bash +cd vime/ +bash scripts/run-deepseek-r1.sh +``` + +On other nodes, you need to join the Ray cluster with the following command: + +```bash +ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" +``` + +Alternatively, if you have a list of all node IPs, for example, an MPI hostfile (where each line is `ip slot=8`), you can add the following commands after the `ray start --head` command in `scripts/run-deepseek-r1.sh`. This allows you to execute the training entirely from node0: + +```bash +for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do + if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 vllm ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" & +done +wait +``` + +### Parameter Introduction + +```bash +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/deepseek-v3.sh" +``` + +This reads the model's config from [scripts/models/deepseek-v3.sh](https://github.com/vllm-project/vime/blob/main/scripts/models/deepseek-v3.sh). These configs are all Megatron parameters. When training with Megatron, it cannot read the model config from the checkpoint, so we need to configure it ourselves. We provide some examples in [scripts/models](https://github.com/vllm-project/vime/tree/main/scripts/models/). + +#### CKPT\_ARGS + +```bash +CKPT_ARGS=( + # HF ckpt required by vllm, we also read the tokenizer from here + --hf-checkpoint $BASE_DIR/DeepSeek-R1/ + #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ + --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/ + # Actor's load directory, if empty, it will read from `ref_load` + --load $BASE_DIR/DeepSeek-R1_vime/ + --save $BASE_DIR/DeepSeek-R1_vime/ + --save-interval 20 +) +``` + +vime will perform online quantization during training based on the quantization configuration in `hf_checkpoint`. For instance, in the current example, we are using the fp8 checkpoint of DeepSeek R1. This means that when updating parameters, we will first perform blockwise quantization on the parameters before passing them to vllm. + +#### PERF\_ARGS + +A set of Megatron parallelism parameters. Only `--use-dynamic-batch-size` and `--max-tokens-per-gpu` are added by vime. + +For the Megatron part, we have configured TP8, PP4, CP4, and EP32. Since DeepSeek-R1 has 61 layers, which is not divisible by 4, we have specifically configured the last pipeline stage to have 13 layers. + +`max_tokens_per_gpu` refers to the maximum number of tokens each GPU can process. When `use_dynamic_batch_size` is enabled, it will pack data of varying lengths within a batch as close to `max_tokens_per_gpu`. If a single data item exceeds `max_tokens_per_gpu`, it will form its own batch without truncation. When context parallelism (CP) is enabled, it allows CP GPUs to share a total length of `CP * max_tokens_per_gpu` tokens. + +When `dynamic_batch_size` is enabled, the traditional `micro_batch_size` is ignored. + +⚠️ vime always trains the model using data packing and strictly guarantees per-sample or per-token loss. This means enabling dynamic batch size will not affect the loss calculation. It is recommended to enable it. + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 13 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) +``` + +#### GRPO\_ARGS + +Currently, these are some GRPO-related parameters in vime: + +```bash +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) +``` + +If you wish to train without loading the reference model, you need to remove `--use-kl-loss` and set `--kl-coef 0.00` (the default value is 0). + +#### OPTIMIZER\_ARGS + +We have configured CPU Adam with the following parameters to save GPU memory. + +```bash +OPTIMIZER_ARGS=( + ... + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) +``` + +#### VLLM\_ARGS + +These are the parameters required by vLLM. `--rollout-num-gpus-per-engine` is the total worker GPU count for one engine; here it is `tensor_parallel_size * data_parallel_size`, not just the tensor-parallel size. Other vLLM parameters are passed to vime by adding the `--vllm-` prefix. To fully leverage vLLM's large EP inference capabilities, we enable `--vllm-enable-expert-parallel` for expert parallelism and `--vllm-data-parallel-size 8` for data-parallel attention. DeepEP is available but disabled by default (see commented flags in the script). + +The final `--vllm-server-concurrency` is a parameter specific to vime. It is used to prevent the vllm server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128. + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.7 + --vllm-enable-expert-parallel + + # dp attention + --vllm-data-parallel-size 8 + + # enable deepep for vllm + + # mtp + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +#### MISC\_ARGS + +Some additional Megatron configurations. Note that Megatron's deepep is configured here. + +```bash +MISC_ARGS=( + ... + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) +``` diff --git a/docs/en/examples/glm4-9B.md b/docs/en/examples/glm4-9B.md new file mode 100644 index 000000000..92d925fb3 --- /dev/null +++ b/docs/en/examples/glm4-9B.md @@ -0,0 +1,278 @@ +# GLM4-9B with 8xH100 + +## Environment Setup + +After pulling the `vllm/vime:latest` image, initialize the image environment as follows: + +```bash +cd /root/ +git clone https://github.com/vllm-project/vime.git +cd vime/ +pip install -e . --no-deps +``` + +Download the model and data: + +```bash +# hf checkpoint +hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 + +# train data +hf download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +# eval data +hf download --repo-type dataset zhuzilin/aime-2024 \ + --local-dir /root/aime-2024 +``` + +Convert the Hugging Face checkpoint to a Megatron-loadable Hugging Face checkpoint: + +```bash +# mcore checkpoint +cd /root/vime +source scripts/models/glm4-9B.sh +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/GLM-Z1-9B-0414 \ + --save /root/GLM-Z1-9B-0414_torch_dist +``` + +## Run Training + +Execute the training: + +```bash +cd /root/vime +bash scripts/run-glm4-9B.sh +``` + +### Parameter Introduction + +Here, we will briefly introduce the various components of the [run-glm4-9B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-glm4-9B.sh) script: + +#### MODEL\_ARGS + +```bash +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/glm4-9B.sh" +``` + +Reads the model's config from [scripts/models/glm4-9B.sh](https://github.com/vllm-project/vime/blob/main/scripts/models/glm4-9B.sh). These configs are all Megatron parameters. When training with Megatron, it cannot read the model config from the checkpoint, so we need to configure it ourselves. We provide some examples in [scripts/models](https://github.com/vllm-project/vime/tree/main/scripts/models/). + +⚠️ Ensure that settings such as `--rotary-base` in the model configuration file match the settings of the model you are currently training. This is because different models, even with the same architecture, might use different values. If needed, you can override these parameters in your script after loading the model weights. For instance: + +```bash +source "${SCRIPT_DIR}/models/glm4-9B.sh" + +MODEL_ARGS += ( --rotary-base 10000 ) +``` + +#### CKPT\_ARGS + +```bash +CKPT_ARGS=( + # HF checkpoint required by vllm; we also read the tokenizer from here + --hf-checkpoint /root/GLM-Z1-9B-0414 + # Checkpoint for the reference model + --ref-load /root/GLM-Z1-9B-0414_torch_dist + # Load directory for the actor; if empty, it will be loaded from `ref_load` + --load /root/GLM-Z1-9B-0414_vime/ + --save /root/GLM-Z1-9B-0414_vime/ + --save-interval 20 +) +``` + +#### ROLLOUT\_ARGS + +```bash +ROLLOUT_ARGS=( + # Prompt dataset, each line is a JSON object + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + # If the `input_key` in the prompt contains an OpenAI message, + # tokenizer.apply_chat_template(...) will be executed + --apply-chat-template + # Whether to shuffle the data + --rollout-shuffle + + # Reward model type. + # vime provides many types and --custom-rm-path for custom models + --rm-type deepscaler + + # Total number of rollouts to train + --num-rollout 3000 + # Number of prompts in one rollout + --rollout-batch-size 32 + # Number of responses to sample per prompt + # A rollout will have rollout_batch_size * n_samples_per_prompt items + --n-samples-per-prompt 8 + # Rollout sampling parameters + --rollout-max-response-len 8192 + --rollout-temperature 1 + + # Number of training steps corresponding to one rollout + --num-steps-per-rollout 1 + # Whether to balance data during training, which might improve speed + --balance-data +) +``` + +#### EVAL\_ARGS + +During evaluation, most rollout parameters are inherited, but we provide some parameters that can override the rollout configuration, allowing for different sampling strategies for training and evaluation. + +```bash +EVAL_ARGS=( + --eval-interval 5 + --eval-prompt-data /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) +``` + +#### PERF\_ARGS + +A set of Megatron's parallelism parameters. Only `--use-dynamic-batch-size` and `--max-tokens-per-gpu` are added by vime. + +`max_tokens_per_gpu` specifies the maximum number of tokens each GPU can process. When `use_dynamic_batch_size` is enabled, it will try to pack data of varying lengths within a batch up to `max_tokens_per_gpu`, thus forming a dynamic micro-batch size. If a single data item's length exceeds `max_tokens_per_gpu`, it will form its own batch without being truncated. When context parallelism (CP) is enabled, it allows the CP GPUs to share data with a total length of `CP * max_tokens_per_gpu` tokens. + +When `dynamic_batch_size` is enabled, the traditional `micro_batch_size` is ignored. + +⚠️ vime always trains the model using data packing and strictly guarantees per-sample or per-token loss. This means enabling dynamic batch size will not affect the loss calculation. It is recommended to enable it. + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 2 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 4608 +) +``` + +#### GRPO\_ARGS + +Here are some GRPO-related parameters: + +```bash +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) +``` + +#### OPTIMIZER\_ARGS + +```bash +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) +``` + +#### VLLM\_ARGS + +These are the parameters required by vLLM. With the default parallel settings, `--rollout-num-gpus-per-engine` corresponds to vLLM's `tensor_parallel_size`. Other vLLM parameters are passed to vime by adding the `--vllm-` prefix. + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 +) +``` + +### Co-located Training and Inference + +In the original script, the resource configuration is as follows: + +```bash +ray job submit ... \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + ... +``` + +This enables decoupled training and inference, where the training part will use 1 machine with 4 GPUs, and the inference will use another 4 GPUs. + +If you want to use the co-located feature, you need to add `--colocate` and remove `--rollout-num-gpus`: + +```bash +ray job submit ... \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ... +``` + +In this case, both training and inference will share these 8 GPUs. + +⚠️ When using co-located training and inference, Megatron will always occupy some GPU memory. Therefore, you need to adjust `--vllm-gpu-memory-utilization` to reduce the proportion of memory occupied by vllm. + +### Dynamic Sampling + +vime supports more complex sampling schemes, such as the dynamic sampling in [DAPO](https://dapo-sia.github.io/). To enable dynamic sampling, you need to configure: + +```bash + --over-sampling-batch-size ${OVER_SAMPLING_BS} \ + --dynamic-sampling-filter-path \ + vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ +``` + +Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example: + +```bash + --rollout-batch-size 32 \ + --n-samples-per-prompt 8 \ + --over-sampling-batch-size 64 \ +``` + +The sampling will then directly sample 64 prompts, with 8 samples per prompt. Since vime performs asynchronous sampling internally, we will receive the 8 responses for each prompt sequentially. Upon receiving responses, they will be filtered using the function specified by `dynamic_sampling_filter_path`. If they pass, these 8 data points are kept; otherwise, they are discarded. The function in the example checks if the answers are all correct or all incorrect: + +```python +def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): + rewards = [sample.reward for sample in samples] + return torch.tensor(rewards, dtype=torch.float).std() > 0.0 +``` + +When we have received 32 \* 8 data points, we will immediately stop sampling and will not wait for the remaining data to be sampled. If more than 32 prompts' worth of data is discarded (leaving fewer than 32 prompts' worth), we will then sample another 64 prompts. + +### Partial Rollout + +During the process of dynamic sampling, a large number of requests are aborted prematurely. We can configure the `--partial-rollout` parameter to save these partially generated requests to a data buffer. In the next rollout, these requests can be retrieved to continue data generation, thereby further optimizing performance. + +You can customize how data is retrieved from the buffer by configuring the `--buffer-filter-path`. The default function is: + +```python +def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]: + num_to_pop = min(len(buffer), num_samples) + samples = buffer[:num_to_pop] + del buffer[:num_to_pop] + return samples +``` + +This means that each time, the data corresponding to the first `num_samples` prompts is retrieved, totaling `num_samples * n_samples_per_prompt` items. + +⚠️ The `sample.metadata` of each partial rollout sample stores the rollout ID from its initial generation, which can be used for data filtering. diff --git a/docs/en/examples/glm4.7-30B-A3B.md b/docs/en/examples/glm4.7-30B-A3B.md new file mode 100644 index 000000000..851bba232 --- /dev/null +++ b/docs/en/examples/glm4.7-30B-A3B.md @@ -0,0 +1,141 @@ +# GLM-4.7-Flash with 8×H100 + +## Environment Preparation + +The environment setup, data, and checkpoint conversion are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7-Flash. + +### Download Model + +```bash +hf download THUDM/GLM-4.7-Flash --local-dir /root/GLM-4.7-Flash +``` +### Convert Checkpoint + +To convert the Hugging Face checkpoint to torch_dist format: + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm4.7-30B-A3B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/GLM-4.7-Flash/ \ + --save /root/GLM-4.7-Flash_torch_dist/ +``` + +## Run Training + +Execute the training script: + +```bash +cd /root/vime +bash scripts/run-glm4.7-30B-A3B.sh +``` + +### Parameter Introduction + +Here, we will briefly introduce the key parts in [run-glm4.7-30B-A3B.sh](../../../scripts/run-glm4.7-30B-A3B.sh). + +#### MoE Configuration + +GLM-4.7-Flash is a Mixture-of-Experts (MoE) model with 64 routed experts (top-4 activation) and 1 shared expert. It has 47 layers: 1 dense layer + 46 MoE layers. + +1. To support running GLM-4.7-Flash on 8×H100, we need to enable Megatron's CPU Adam to save GPU memory: + + ```bash + OPTIMIZER_ARGS=( + ... + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + ) + ``` + +2. Enable MoE optimization in Megatron. For single-node 8×H100, we use TP=1, EP=8: + + ```bash + PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + ... + ) + ``` + +3. Enable vLLM data parallelism for MoE: + + ```bash + VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 8 + --vllm-enable-expert-parallel + ... + ) + ``` + +#### MTP Speculative Decoding (Inference Acceleration) + +GLM-4.7-Flash includes 1 MTP (Multi-Token Prediction) layer, which can be used for speculative decoding during inference to speed up rollout generation. To enable this, add the following to `VLLM_ARGS`: + +```bash +VLLM_ARGS=( + ... + # MTP speculative decoding + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +This enables vLLM to use the model's MTP layer for speculative decoding. The MTP layer predicts multiple future tokens, and vLLM verifies them in parallel, leading to faster generation. + +> ⚠️ **Note**: Speculative decoding requires additional GPU memory. If you encounter OOM issues, try reducing `--vllm-gpu-memory-utilization` or disabling speculative decoding. + +#### MTP Training + +vime also supports training MTP layers jointly with the main model for models that have MTP weight conversion implemented (e.g., MiMo, GLM-4.7). When enabled, the relevant arguments are: + +```bash +# Add MTP layer count to model config +MODEL_ARGS+=(--mtp-num-layers 1) + +# Enable MTP training +SPEC_ARGS=( + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) +``` + +- `--mtp-num-layers 1`: Tells Megatron to load the MTP layer from the checkpoint. +- `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen. +- `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2. + +> **Note**: The native DeepSeek-layout loader uses the model's configured layer count when mapping MTP weights, including GLM-4.7-Flash's 47-layer architecture. +> +> For other models with MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference. + +### Multi-Node Adaptation + +The checked-in `scripts/run-glm4.7-30B-A3B.sh` launcher starts a local, single-node Ray cluster and passes `--actor-num-nodes 1`; it is not a drop-in multi-node launcher. To adapt this recipe for multi-node training (for example, 2×8 H100), start or connect all workers to the same Ray cluster and update the launcher as follows: + +Key modifications for multi-node: + + - Place the model and data on a path accessible by all nodes. + - Set `MASTER_ADDR` to an address accessible by all nodes. + - Set `--actor-num-nodes` to the number of training nodes instead of `1`. + - Remove CPU Adam configurations (distributed optimizer reduces per-GPU memory usage). + - Adjust parallelism: e.g., TP=4, PP=2, EP=8, CP=2. + +When the total number of GPUs is not a multiple or divisor of the total number of experts (64), you can use `--vllm-eplb-config` to add redundant experts. For example, in a 24-GPU scenario: + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 24 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 3 + --vllm-enable-expert-parallel + --vllm-eplb-config '{"num_redundant_experts": 16}' +) +``` diff --git a/docs/en/examples/glm4.7-355B-A32B.md b/docs/en/examples/glm4.7-355B-A32B.md new file mode 100644 index 000000000..14580fc18 --- /dev/null +++ b/docs/en/examples/glm4.7-355B-A32B.md @@ -0,0 +1,174 @@ +# GLM-4.7 with 64xH100 + +## Environment Preparation + +The environment setup and dataset download are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7. + +### Prerequisites + +GLM-4.7 follows the standard vime Docker environment. For multi-node launches, make sure all nodes can access the same `$BASE_DIR` path and unset proxy variables before starting Ray workers: + +```bash +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY +``` + +### Download Model + +```bash +hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B +``` + +### Convert Checkpoint + +To convert the Hugging Face checkpoint to torch_dist format, use 2 nodes x 8 GPUs: + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm4.5-355B-A32B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=2 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \ + --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/ +``` + +Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` is the node index, configured just like a multi-node `torchrun` job. + +## Run Training + +Execute the training script from node0: + +```bash +cd /root/vime +export BASE_DIR=/shared/path # accessible by all nodes +bash scripts/run-glm4.7-355B-A32B.sh +``` + +### Parameter Introduction + +Here, we briefly introduce the key parts in the [run-glm4.7-355B-A32B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-glm4.7-355B-A32B.sh) script. + +#### MoE Configuration + +GLM-4.7 is a Mixture-of-Experts (MoE) model with 160 routed experts (top-8 activation) and shared experts. It has 92 layers: 3 dense layers + 89 MoE layers. + +1. To support GLM-4.7 on 64xH100, we enable Megatron's CPU Adam to save GPU memory: + + ```bash + OPTIMIZER_ARGS=( + ... + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + ) + ``` + +2. Enable MoE optimization in Megatron. For the provided 64xH100 example, we use TP=8, PP=4, CP=2, and EP=16: + + ```bash + PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 2 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + ... + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 + ) + ``` + +3. Enable MoE optimization in vLLM: + + ```bash + VLLM_ARGS=( + --rollout-num-gpus-per-engine 32 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 4 + --vllm-enable-expert-parallel + ... + ) + ``` + +#### MTP Speculative Decoding (Inference Acceleration) + +GLM-4.7 includes MTP (Multi-Token Prediction) layers that can be used for speculative decoding during inference to speed up rollout generation. To enable this, add the following to `VLLM_ARGS`: + +```bash +VLLM_ARGS=( + ... + # MTP speculative decoding + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +This lets vLLM use the model's MTP layer as the draft model for MTP speculative decoding. + +> ⚠️ **Note**: Speculative decoding requires additional GPU memory. If you encounter OOM issues, try reducing `--vllm-gpu-memory-utilization` or disabling speculative decoding. + +#### MTP Training + +vime also supports training the MTP layers jointly with the main model for GLM-4.7. When enabled, the relevant arguments are: + +```bash +# Add MTP layer count to model config +MODEL_ARGS+=(--mtp-num-layers 1) + +# Enable MTP training +MTP_ARGS=( + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) +``` + +- `--mtp-num-layers 1`: Tells Megatron to load the MTP layer from the checkpoint. +- `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen. +- `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2. + +> **Note**: The native loader in `vime/backends/megatron_utils/hf_to_megatron/glm.py` maps both regular and MTP weights. + +#### Multi-Node Support + +This example already targets multi-node training. Before launching: + +- Place the model checkpoints and datasets on a path accessible by all nodes. +- Set `MASTER_ADDR` to an address reachable by all nodes. +- Unset proxy variables before starting Ray workers. +- Provide a `HOSTFILE` listing worker IPs (one per line) and export `HOSTFILE=/path/to/hostfile` before launching. +- Adjust parallelism coherently. The default example uses TP=8, PP=4, EP=16, CP=2, while rollout uses 32 GPUs per engine with vLLM DP attention. + +If your rollout GPU count does not divide the expert count cleanly, you can use `--vllm-eplb-config` to enable EPLB and configure redundant experts. + +## FP8 Rollout + +The open-source FP8 checkpoint of GLM-4.7 uses per-channel quantization, which cannot currently enable DeepEP in vLLM. You can convert it to a 128x128 per-block FP8 checkpoint with the tool provided in vime: + +```bash +cd /root/vime +python tools/convert_hf_to_fp8.py \ + --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \ + --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \ + --strategy block --block-size 128 128 \ + --max-workers 4 +``` + +Then switch `--hf-checkpoint` to `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` to enable FP8 rollout. + +An example FP8 `VLLM_ARGS` setup is: + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 32 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 32 + --vllm-enable-expert-parallel + --vllm-cudagraph-capture-sizes 5 10 20 40 $(seq 80 40 640) + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' + --vllm-all2all-backend deepep_high_throughput +) +``` diff --git a/docs/en/examples/glm5.2-744B-A40B.md b/docs/en/examples/glm5.2-744B-A40B.md new file mode 100644 index 000000000..76ffffd29 --- /dev/null +++ b/docs/en/examples/glm5.2-744B-A40B.md @@ -0,0 +1,168 @@ +# GLM-5.2 744B-A40B with 256xH100 + +This is the recommended 32-node, 256-H100 training example for [GLM-5.2](https://z.ai/blog/glm-5.2). + +The recipe uses the GLM-5.2 BF16 checkpoint for Megatron training and the FP8 checkpoint for vLLM rollout. It assumes two Hugging Face repositories will be available: + +- BF16: `zai-org/GLM-5.2` +- FP8: `zai-org/GLM-5.2-FP8` + +## Environment Setup + +For environment setup and dataset download, see [Example: Qwen3-4B](qwen3-4B.md). For multi-node training, make sure every node can access the same `$BASE_DIR` path. + +### Download Model + +```bash +hf download zai-org/GLM-5.2 --local-dir $BASE_DIR/GLM-5.2 +hf download zai-org/GLM-5.2-FP8 --local-dir $BASE_DIR/GLM-5.2-FP8 +``` + +The open-source GLM-5.2 config uses `model_type: glm_moe_dsa`, which vime maps onto +the native DeepSeek-V3.2 loader since the two share the +same DSA weight layout. + +### Convert Checkpoint + +The training side needs the BF16 Hugging Face checkpoint converted to the Megatron torch_dist format. The torch_dist format is reshardable, so the conversion parallel layout does **not** need to match training; we use a layout that satisfies Megatron's expert-group constraint on the conversion node count. + +Run the following on 4 nodes / 32 GPUs: + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm5.2-744B-A40B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=4 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --tensor-model-parallel-size 8 \ + --pipeline-model-parallel-size 2 \ + --decoder-last-pipeline-num-layers 40 \ + --expert-model-parallel-size 16 \ + --expert-tensor-parallel-size 1 \ + --hf-checkpoint $BASE_DIR/GLM-5.2/ \ + --save $BASE_DIR/GLM-5.2_torch_dist/ +``` + +Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` is the current node index. + +`MODEL_ARGS` includes `--allgather-cp`, a vime-only flag, so `tools/convert_hf_to_torch_dist.py` registers it too (it is a no-op for conversion). On 32 GPUs, Megatron requires `expert_tp(1) * expert_model_parallel * pp` to divide the world size, so we convert with `EP=16` (`1*16*2=32`). The resulting checkpoint still loads at training-time `EP=32` because torch_dist is reshardable. + +## Run Training + +From node0: + +```bash +cd /root/vime +export BASE_DIR=/shared/path +export MASTER_ADDR= +export HOSTFILE=$BASE_DIR/hostfile # one worker IP per line, all 32 nodes +bash scripts/run-glm5.2-744B-A40B.sh +``` + +If `HOSTFILE` is not set, join the other nodes to the Ray cluster manually. + +### Parameter Introduction + +#### Model Configuration + +`scripts/models/glm5.2-744B-A40B.sh` contains the GLM-5.2 DSA + cross-layer index sharing configuration: 256 routed experts, top-8 activation, 1 shared expert, and 78 layers total (3 dense + 75 MoE). + +The DSA index sharing schedule, such as `index_topk_freq=4` and `index_skip_topk_offset=3`, is read from the Hugging Face config. The Megatron side uses the shared `vime_plugins.models.glm5.glm5:get_glm5_spec` provider and enables: + +```bash +--allgather-cp +``` + +This makes DSA + context parallel use the allgather-CP layout, and the index-share provider gathers index K/V across the CP group. + +#### Training Parallelism + +The default script targets 32 nodes and 256 GPUs: + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 4 + --pipeline-model-parallel-size 8 + --decoder-first-pipeline-num-layers 14 + --decoder-last-pipeline-num-layers 16 + --context-parallel-size 8 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + ... +) +``` + +`TP=4 * PP=8 * CP=8 = 256` GPUs form one training group (`DP=1`). The expert group constraint `expert_tp(1) * EP(32) * PP(8) = 256` divides the world size exactly (`expert_dp=1`). + +DSA cross-layer index sharing requires every pipeline stage to **start** on a "computing" layer. With `index_topk_freq=4` / `index_skip_topk_offset=3`, the computing layers are 1, 2, 3, 7, 11, ..., 75. A uniform `78/8` split would start stages on skip layers and fail the index-share assertion in `get_glm5_spec`. We therefore use `--decoder-first-pipeline-num-layers 14` and `--decoder-last-pipeline-num-layers 16`, leaving 6 middle stages of `(78-14-16)/6 = 8` layers each. The stage starts land on global layers 1, 15, 23, 31, 39, 47, 55, 63 — all computing layers. + +#### BF16 Training + FP8 Rollout + +The launcher writes the default paths directly in `CKPT_ARGS` and `ROLLOUT_ARGS`, matching the style of the other example scripts: + +```bash +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/GLM-5.2-FP8 + --ref-load $BASE_DIR/GLM-5.2_torch_dist + --load $BASE_DIR/GLM-5.2_vime + --save $BASE_DIR/GLM-5.2_vime + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + ... +) +``` + +`--hf-checkpoint` provides FP8 weights and the tokenizer for vLLM rollout; `--ref-load` is the Megatron torch_dist checkpoint converted from BF16. To debug BF16 rollout, change `--hf-checkpoint` in the script to `$BASE_DIR/GLM-5.2`. + +#### vLLM Configuration + +The rollout side runs with **prefill/decode (PD) disaggregation**: 1 prefill engine (64 GPU) + 3 decode engines (192 GPU) = 256 GPUs total (which must equal the colocated `rollout_num_gpus`). Each engine spans 64 GPUs with data parallelism and vLLM expert parallelism. Prefill uses the high-throughput DeepEP backend; decode uses the low-latency backend. The split is configured via the `--vllm-config` YAML: + +```yaml +vllm: + - name: default + server_groups: + - worker_type: prefill + num_gpus: 64 + num_gpus_per_engine: 64 + overrides: { data_parallel_size: 64, enable_expert_parallel: true, all2all_backend: deepep_high_throughput, kv_transfer_config: { kv_connector: MooncakeConnector, kv_role: kv_producer, ... }, ... } + - worker_type: decode + num_gpus: 192 + num_gpus_per_engine: 64 + overrides: { data_parallel_size: 64, enable_expert_parallel: true, max_cudagraph_capture_size: 72, all2all_backend: deepep_low_latency, kv_transfer_config: { kv_connector: MooncakeConnector, kv_role: kv_consumer, ... }, ... } +``` + +The source `mooncake` transport maps to vLLM's `MooncakeConnector`. The source IB-device list maps to `kv_connector_extra_config.device_name`; the prefill and decode groups use `kv_producer` and `kv_consumer`, respectively. + +The shared rollout arguments use vLLM-native FP8 KV cache and CUDA-graph settings: + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.70 + --vllm-kv-cache-dtype fp8_e4m3 + --vllm-max-cudagraph-capture-size 48 + --vllm-config "${VLLM_CONFIG_FILE}" +) +``` + +MTP / EAGLE speculative decoding is enabled using the model's own next-token-prediction layer (the GLM-5.2 checkpoint ships an MTP layer), so no separate draft model is needed: + +```bash +--vllm-speculative-config '{"method":"mtp","num_speculative_tokens":5}' +``` + +vLLM measures CUDA-graph capture size in flattened query tokens. With five speculative tokens, each decode request contributes `1 + 5 = 6` query tokens. The shared limit `48` therefore covers 8 requests, while the decode-group override `72` covers 12 requests. vLLM derives the DeepEP dispatch-buffer size from its scheduler token capacity. + +`VLLM_ENGINE_ITERATION_TIMEOUT_S=3600` raises vLLM's engine watchdog for this long-running multi-node workload. + +#### Networking + +DeepEP/NVSHMEM communication across nodes needs the IB-aware NCCL settings in the Ray runtime env (`NCCL_SOCKET_IFNAME`, `NCCL_IB_*`, `NCCL_NET_GDR_LEVEL`, `NCCL_P2P_LEVEL=NVL`, `NCCL_NVLS_ENABLE=0`, `MC_IB_PCI_RELAXED_ORDERING`, ...). The script defaults to `SOCKET_IFNAME=eth0`; set `SOCKET_IFNAME` before launch if your environment differs, and it will be written to `GLOO_SOCKET_IFNAME`, `TP_SOCKET_IFNAME`, and `NCCL_SOCKET_IFNAME`. DeepEP also requires `NVSHMEM_DISABLE_NCCL=1`. diff --git a/docs/en/examples/qwen3-30B-A3B.md b/docs/en/examples/qwen3-30B-A3B.md index bca56238d..0d266debc 100644 --- a/docs/en/examples/qwen3-30B-A3B.md +++ b/docs/en/examples/qwen3-30B-A3B.md @@ -31,7 +31,7 @@ bash scripts/run-qwen3-30B-A3B.sh Here, we will briefly introduce the MoE-related parts in the [run-qwen3-30B-A3B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-30B-A3B.sh) script. -1. To support running Qwen3-30B-A3B in an 8xH800 environment, we need to enable Megatron's CPU Adam to save GPU memory. The corresponding configuration is: +1. To support running Qwen3-30B-A3B in an 8xH100 environment, we need to enable Megatron's CPU Adam to save GPU memory. The corresponding configuration is: ```bash OPTIMIZER_ARGS=( @@ -72,92 +72,66 @@ Here, we will briefly introduce the MoE-related parts in the [run-qwen3-30B-A3B. For DP on the attention block plus EP on the experts, combine `--vllm-data-parallel-size N` with `--vllm-enable-expert-parallel`. -### Multi-Node Support - -The following uses **two machines with 8 GPUs each (16 GPUs total)** as the starting example; scripts and parameters scale to **N nodes**. Key differences from single-node: - -- Place weights, checkpoints, and data on storage visible to every node (e.g. NFS). -- Set `MASTER_ADDR` to the head **LAN IP** (not `127.0.0.1`). -- Omit CPU Adam (multi-node uses a distributed optimizer; do not use `--optimizer-cpu-offload`). -- `global-batch-size` must equal `rollout-batch-size × n-samples-per-prompt`. +### BF16 Training with FP8 Inference -#### Topology +vime also supports BF16 training with FP8 inference. For the Qwen3-30B-A3B model, just download the FP8 weights: -| Component | Dual-node defaults | -|-----------|-------------------| -| Cluster | `ACTOR_NUM_NODES=2`, `ACTOR_NUM_GPUS_PER_NODE=8` | -| Megatron training | TP=8, EP=8, CP=2 (experts sharded across nodes) | -| vLLM rollout | Cross-node TP=16 (`rollout-num-gpus-per-engine = nodes × GPUs per node`) | -| Scheduling | Ray cluster + `--colocate` mode | - -Convert checkpoints with Megatron parallelism matching training (dual-node: TP=8, EP=8). Checkpoint EP must match `--expert-model-parallel-size`, or `load_checkpoint` may hang or resharding may be extremely slow. - -#### Start the Ray Cluster +```bash +hf download Qwen/Qwen3-30B-A3B-FP8 --local-dir /root/Qwen3-30B-A3B-FP8 +``` -Start Ray **outside** the training script on each node. Join all workers first; verify `ray status` reports the expected GPU count, then submit training from the head. Dual-node example: +And replace `--hf-checkpoint` in the script with: ```bash -# === Head node === -export MASTER_ADDR= -ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ - --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# === Each worker node === -export MASTER_ADDR= -ray start --address="${MASTER_ADDR}:6379" --node-ip-address= --num-gpus 8 +#--hf-checkpoint /root/Qwen3-30B-A3B +--hf-checkpoint /root/Qwen3-30B-A3B-FP8 ``` -See [Quick Start — Multi-node training](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models) for more details. - -#### Run Training +This triggers FP8 inference. Currently we directly cast the BF16 weights to FP8; more precision-friendly quantization schemes will be added over time. -After the Ray cluster is ready, on the **head node** set multi-node env vars and run the **same script as single-node** (`ACTOR_NUM_NODES>1` skips Ray startup and applies multi-node defaults): +⚠️ The Megatron checkpoint used for training must still be the one originally converted from the BF16 huggingface weights (`--ref-load` / `--load` unchanged). -```bash -export MASTER_ADDR= -export ACTOR_NUM_NODES=2 -export ACTOR_NUM_GPUS_PER_NODE=8 -cd /root/vime -bash scripts/run-qwen3-30B-A3B.sh -``` +### Multi-Node Support -2-step smoke test: +The following uses **2 nodes × 8 GPUs (16 GPUs total) in colocate mode** as the example. The only differences from single-node are "starting Ray across nodes" and "adjusting a few resource/parallelism arguments"; the training script itself is unchanged. -```bash -NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh -``` +1. **Shared storage**: put the model, data, and checkpoints on a location that every node can access at the same path (e.g. NFS). -To scale to N nodes (e.g. 4×8), join all workers to Ray, set `ACTOR_NUM_NODES=4` on the head, and tune `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE` for total GPU count. +2. **Start Ray across nodes** (outside the training script, run manually on each node; see [Quick Start — Multi-node training](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models)): -#### Key Multi-Node Parameters + ```bash + # Head node (node0); MASTER_ADDR must be a LAN IP, not 127.0.0.1 + export MASTER_ADDR= + ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats -| Variable | Dual-node default | Description | -|----------|-------------------|-------------| -| `ACTOR_NUM_NODES` | 2 (default 1 for single-node) | Total nodes including head; script skips Ray startup when >1 | -| `ACTOR_NUM_GPUS_PER_NODE` | 8 | GPUs per node | -| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron parallelism | -| `ROLLOUT_NUM_GPUS_PER_ENGINE` | total GPUs | vLLM engine GPU count | -| `ENABLE_R3` | 1 | set to 0 to disable R3 | + # Every other node + ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 + ``` -Default batch: `rollout-batch-size=4`, `n-samples-per-prompt=2`, `global-batch-size=8`; vLLM uses `--vllm-moe-backend triton`. + Wait until `ray status` reports 16 GPUs before submitting. Because you started the cluster manually, make the script skip its process-management preamble — remove (or comment out) **both** the initial cleanup block (`ray stop --force`, `pkill -9 ray`, `pkill -9 python`, `pkill -9 redis`) **and** the `ray start --head ...` line. Otherwise running the script tears down the head you just started (and orphans the workers), so `ray job submit` to `http://127.0.0.1:8265` fails. Keep the rest of the script — it still sources the model args and runs `ray job submit`. -#### Multi-Node Troubleshooting +3. **Adjust script arguments** (`scripts/run-qwen3-30B-A3B.sh`): + - Change `--actor-num-nodes` for `train.py` from `1` to `2` (keep `--actor-num-gpus-per-node` at 8). Under colocate, `--rollout-num-gpus` is auto-set to `actor_num_gpus_per_node × actor_num_nodes = 16`, so you don't set it manually. + - Scale up the parallelism in `PERF_ARGS` for the doubled GPU count (e.g. raise TP or add DP); for concrete large-scale ratios see the bigger-cluster examples such as [GLM-4.7](glm4.7-355B-A32B.md) and [DeepSeek-R1](deepseek-r1.md). + - `global-batch-size` must equal `rollout-batch-size × n-samples-per-prompt`. + - (Optional) Multi-node uses a distributed optimizer, which lowers optimizer memory pressure, so you may drop the CPU Adam options (`--optimizer-cpu-offload`, etc.) from `OPTIMIZER_ARGS` for speed. -- **Worker cannot join Ray / NCCL failures**: check `MASTER_ADDR`, container `/etc/hosts` (hostname must not map to `127.0.0.1`), `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`. -- **`Not enough samples X for global_batch_size Y`**: keep `global-batch-size` equal to `rollout-batch-size × n-samples-per-prompt`. -- **GPU memory full but no processes**: restart the container or run `ray stop --force` to clear stale vLLM contexts. +4. **Keep each vLLM engine within a single node**: prefer `--rollout-num-gpus-per-engine 8` (one engine per node) over `16` (a single engine spanning both nodes at TP=16). Cross-node TP is noticeably slower and more sensitive to per-token numerics; this value must divide the total rollout GPU count (16 here). -#### EPLB +⚠️ Common issues: +- **Worker cannot join Ray / NCCL failures**: check `MASTER_ADDR`, container `/etc/hosts` (hostname must not map to `127.0.0.1`), and set `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` on multi-NIC hosts. +- **`Not enough samples X for global_batch_size Y`**: keep `global-batch-size = rollout-batch-size × n-samples-per-prompt`. +- **Fewer than 8 GPUs per node (colocate)**: set `--num-gpus-per-node` explicitly. -When the total number of GPUs is not a multiple or divisor of the total number of experts, enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`. For example, in a 24-GPU scenario: +In addition, when the total number of GPUs is not a multiple or divisor of the total number of experts, you can enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`. For example, in a 24-GPU scenario: - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 24 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 3 - --vllm-enable-expert-parallel - --vllm-enable-eplb - --vllm-eplb-config '{"num_redundant_experts": 16}' - ) - ``` +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 24 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 3 + --vllm-enable-expert-parallel + --vllm-enable-eplb + --vllm-eplb-config '{"num_redundant_experts": 16}' +) +``` diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index ddb6151b2..b9c3815de 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -2,7 +2,7 @@ ## Environment Setup -After pulling the `inferactinc/public:vime-latest` image, initialize the image environment as follows: +After pulling the `vllm/vime:latest` image, initialize the image environment as follows: ```bash cd /root/ @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM\_ARGS -Parameters for vLLM inference. vime uses vLLM as the rollout backend by default (`rollout.py` launches `VLLMEngine`; the default rollout function is `vime.rollout.vllm_rollout.generate_rollout`), so no extra backend flag is needed. `--rollout-num-gpus-per-engine` corresponds to each vLLM engine's `tensor_parallel_size`. Other vLLM parameters are passed to vime with a `--vllm-` prefix (for example, `--vllm-max-model-len`). +These are the parameters required by vLLM. With the default parallel settings, `--rollout-num-gpus-per-engine` corresponds to vLLM's `tensor_parallel_size`. Other vLLM parameters are passed to vime by adding the `--vllm-` prefix. ```bash VLLM_ARGS=( @@ -202,9 +202,7 @@ VLLM_ARGS=( ) ``` -When rollout concurrency is high, tune the vLLM scheduler via the `--vllm-` prefix—for example, `--vllm-max-num-seqs` and `--vllm-max-num-batched-tokens`. Add `--vllm-enforce-eager` for debugging or to work around CUDA graph limits. - -⚠️ vime uses the vLLM router to schedule multiple vLLM servers. With co-located training and inference (`--colocate`), weights are synchronized via CUDA IPC; with decoupled training and inference, the trainer synchronizes weights with vLLM engines over NCCL. +⚠️ vime uses vllm-router to schedule multiple vLLM servers. ### Dynamic Sampling @@ -278,22 +276,21 @@ ray job submit ... \ ... ``` -In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. Like `--actor-num-gpus-per-node`, `--rollout-num-gpus` is a **Ray resource argument** passed to `train.py`: the framework uses it to build the placement group and assign the first bundles to training actors and the remaining bundles to rollout engines (see `vime/ray/placement_group.py`). **Under co-located mode (`--colocate`), this argument is ignored** and is set automatically to `actor_num_gpus_per_node * actor_num_nodes`. Do not put `--rollout-num-gpus` in `VLLM_ARGS`. +In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. -For decoupled training and inference, `VLLM_ARGS` only needs inference-backend settings, for example: +⚠️ If concurrency on each vLLM server is too high, it may exceed the configured CUDA graph capture sizes and affect inference speed. You can adjust this in the following two ways: -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.9 - --vllm-max-num-seqs 256 - --vllm-max-num-batched-tokens 8192 -) -``` +1. Use `--vllm-server-concurrency` to limit the maximum number of concurrent requests sent to one vLLM server. For example: + + ```bash + --vllm-server-concurrency 160 + ``` -Add `--vllm-enforce-eager` when debugging or to work around CUDA graph limits. +2. Use `--vllm-cudagraph-capture-sizes` to configure the CUDA graph sizes initialized by vLLM. For example: -⚠️ When using co-located training and inference, Megatron will always occupy some GPU memory. Reduce vLLM's memory footprint with `--vllm-gpu-memory-utilization`, and reserve headroom for training with `--train-memory-margin-bytes`. + ```bash + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + ``` ### Asynchronous Training diff --git a/docs/en/examples/qwen3-4b-base-openhermes.md b/docs/en/examples/qwen3-4b-base-openhermes.md new file mode 100644 index 000000000..5dc89a278 --- /dev/null +++ b/docs/en/examples/qwen3-4b-base-openhermes.md @@ -0,0 +1,85 @@ +# SFT Qwen3-4B-Base + +## Environment Preparation + +First, we need to create a mirror environment and convert the `Qwen3-4B-Base` model by following the [Example: Qwen3-4B Model](qwen3-4B.md). + +After that, we will process the SFT data. Here, we use the classic [OpenHermes-2.5](https://huggingface.co/datasets/teknium/OpenHermes-2.5) as an example. First, we process the data into a format suitable for `vime` to load. You can use the following script to add a column that conforms to the OpenAI message format and save it to `/root/openhermes2_5.parquet`. + +```python +from datasets import load_dataset + +ds = load_dataset("teknium/OpenHermes-2.5")["train"] + +def convert(sample): + conversations = sample["conversations"] + + def convert_role(role): + if role == "human": + return "user" + elif role == "gpt": + return "assistant" + elif role == "system": + return "system" + else: + raise ValueError(f"Unknown role: {role}") + + messages = [ + { + "role": convert_role(turn["from"]), + "content": turn["value"], + } + for turn in conversations + ] + + return {"messages": messages} + +ds = ds.map(convert) +ds.to_parquet("/root/openhermes2_5.parquet") +``` + +## Execute Training + +Execute the training: + +```bash +cd /root/vime +bash scripts/run-qwen3-4B-base-sft.sh +``` + +### Parameter Introduction + +You can compare [run-qwen3-4B-base-sft.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-4B-base-sft.sh) with [run-qwen3-4B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-4B.sh). You will find that besides changing the model from the instruct version to the base model, the main adjustments are as follows: + +1. Removed `VLLM_ARGS` and `GRPO_ARGS`. This is because it is not necessary to start vLLM or configure GRPO-related settings during the SFT process. + +2. Renamed `ROLLOUT_ARGS` to `SFT_ARGS` and configured it as follows: + + ```bash + SFT_ARGS=( + --rollout-function-path vime.rollout.sft_rollout.generate_rollout + --prompt-data /root/openhermes2_5.parquet + --input-key messages + --rollout-shuffle + --num-epoch 3 + --rollout-batch-size 128 + --global-batch-size 128 + + --loss-type sft_loss + --calculate-per-token-loss + --disable-compute-advantages-and-returns + --debug-train-only + ) + ``` + + SFT actually reuses the custom rollout functionality of vime. By using `--rollout-function-path`, the data generation part is switched from the RL rollout that uses `vllm` to the SFT version that reads data from a file, which is `vime.rollout.sft_rollout.generate_rollout`. + + For SFT, it is recommended to set `rollout_batch_size` and `global_batch_size` to the same value and not to configure `n_samples_per_prompt`. This is equivalent to training one batch right after reading one batch. + + `vime` also supports different loss types, and we configure the SFT loss using `--loss-type sft_loss`. + + As for `--calculate-per-token-loss`, this is because `vime` defaults to calculating the per-sample mean for GRPO. In general SFT training, the average is taken over all unmasked tokens in a batch, so it is recommended to configure this. + + Finally, `--disable-compute-advantages-and-returns` indicates that there is no need to pre-calculate log probabilities during the SFT process, and `--debug-train-only` means that `vllm` does not need to be initialized. + +3. Used `train_async.py` instead of `train.py`. This is to leverage the asynchronous training process to implement data prefetching. diff --git a/docs/en/get_started/NPU.md b/docs/en/get_started/NPU.md index 6a12457fa..b6084fd3e 100644 --- a/docs/en/get_started/NPU.md +++ b/docs/en/get_started/NPU.md @@ -90,7 +90,7 @@ hf download --repo-type dataset zhuzilin/dapo-math-17k \ We provide an example to run GRPO training with [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) on 8 NPUs (4 for the actor, 4 for rollout), please refer to: -[scripts/models/qwen3-4B_npu.sh](https://github.com/vllm-project/vime/blob/npu/scripts/models/qwen3-4B_npu.sh). +[scripts/run-qwen3-4B-npu.sh](../../../scripts/run-qwen3-4B-npu.sh). Just run: ```bash @@ -100,7 +100,8 @@ cd /root/vime source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh -MODEL_ROOT=/root bash scripts/models/qwen3-4B_npu.sh +DATA_ROOT="${MODEL_ROOT:-/root}" bash scripts/run-qwen3-4B-npu.sh \ + 2>&1 | tee /root/vime/train_qwen3_4b_vllm.log ``` The full log is written to `/root/vime/train_qwen3_4b_vllm.log`. diff --git a/docs/en/get_started/agent.md b/docs/en/get_started/agent.md new file mode 100644 index 000000000..e4e01c536 --- /dev/null +++ b/docs/en/get_started/agent.md @@ -0,0 +1,77 @@ +# Agentic RL Training Roadmap + +vime is not limited to single-turn RL. Its main advantage for agentic training is the combination of high-performance training, vLLM rollout serving, and pluggable data-generation interfaces. This makes it suitable for multi-turn tool use, sandbox interaction, subagent branches, context compaction, and test-based rewards. + +This page is a roadmap: use it to decide which docs and examples to read when plugging an agent workflow into vime. + +## Where To Start + +| Goal | Recommended entry point | +| :--- | :--- | +| Run a custom agent loop, tool calls, RAG, browser/terminal/sandbox interaction for each sample | [`--custom-generate-function-path`](customization.md#2-custom-generate-function---custom-generate-function-path), [writing a custom generation function](quick_start.md#writing-custom-generation-function) | +| Implement verifier rewards, test-based rewards, environment success checks, or an external reward service | [`--custom-rm-path`](customization.md#3-reward-model---custom-rm-path), [writing a custom reward function](quick_start.md#writing-custom-reward-function) | +| Return multiple training samples from one prompt, such as subagent, multi-agent, or context-compaction segments | [fan-out return from custom generate](customization.md#returning-multiple-training-samples-for-one-prompt), [`examples/multi_agent`](../_examples_synced/multi_agent/README.md) | +| Avoid blocking training on long-tail agent rollouts | [`examples/fully_async`](../_examples_synced/fully_async/README.md) | +| Study a full end-to-end agent example with sandboxing, real code edits, and test-based grading | [`examples/coding_agent_rl`](../_examples_synced/coding_agent_rl/README.md) | +| Improve vLLM serving throughput for multi-turn agents | [PD Disaggregation](../advanced/pd-disaggregation.md), [vLLM Config](../advanced/vllm-config.md) | +| Enable vLLM optimization flags, router policies, or multi-model serving | [How to Use vLLM](usage.md#how-to-use-vllm), [vLLM Config](../advanced/vllm-config.md), [Speculative Decoding](../advanced/speculative-decoding.md), [Low Precision Training](../advanced/low-precision.md) | + +## Recommended Integration Pattern + +Most agentic RL tasks should start with `--custom-generate-function-path`. This function converts one agent execution into vime-trainable `Sample` objects: fill `tokens`, `response_length`, `loss_mask`, and `status`, then either fill `reward` directly or let `--custom-rm-path` compute it. + +The agent workflow itself may speak in strings, chat messages, tool calls, environment observations, or framework-specific events. The training target, however, should stay token based. Preserve the model-sampled token ids and use `loss_mask` to separate trainable model output from prompt, template, tool-observation, or environment text. + +If one prompt rollout corresponds to one training sample, return a single `Sample`. If one rollout splits into multiple trainable segments, such as subagent trajectories, main-agent continuations, or pre/post-compaction segments, return `list[Sample]` and set the same `rollout_id` on all sibling samples. vime then keeps those samples together for train-step splitting and loss aggregation instead of counting them as independent rollouts. + +Reach for `--rollout-function-path` only when you need to replace the whole rollout orchestration. Common reasons include custom data-source scheduling, cross-rollout background queues, fully asynchronous generation, or workflows that cannot fit the default `vllm_rollout` prompt-by-sample structure. + +## Agent Runtime Adapters + +vime includes protocol adapters for existing agent runtimes: + +- `vime.agent.adapters.AnthropicAdapter`: Anthropic Messages API, used by Claude Code style agents. +- `vime.agent.adapters.OpenAIAdapter`: OpenAI Chat Completions and Responses APIs, used by OpenAI SDK / OpenAI Agents SDK style clients. + +Adapters are a convenience layer, not a separate agent framework. Their contract is message history in, sampled tokens out: they render the chat template, call vLLM with `input_ids` and `return_logprob=True`, and export the returned token ids/logprobs as trainable trajectory segments. They avoid re-tokenizing response text to recover the training target. + +Instantiate the protocol-specific adapter in your custom generate function, run its `app` with aiohttp, then manage each rollout through the adapter instance: + +```python +from vime.agent.adapters import AnthropicAdapter + +adapter = AnthropicAdapter( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, +) + +adapter.open_session(session_id, sampling_defaults=sampling_params) +# Agent client sends requests to adapter.app. +segments = await adapter.finish_session(session_id) +``` + +For multi-turn agents, use a stable `session_id`. The adapters pass it as `X-SMG-Routing-Key` so vLLM can route one session to the same worker and reuse prefix cache. + +## Sandbox Backends + +The coding-agent example ships `vime.agent.sandbox.E2BSandbox` for E2B-compatible services. The rest of the agent lifecycle depends only on the provider-neutral `vime.agent.sandbox.Sandbox` protocol, so Docker, Modal, or local VM backends can implement the same contract without changing rollout logic. See [Porting to a New Sandbox Backend](../_examples_synced/coding_agent_rl/README.md#porting-to-a-new-sandbox-backend). + +## Agent Serving And Performance + +Agentic rollouts tend to depend more heavily on serving configuration than ordinary single-turn generation: contexts are longer, requests are multi-turn, latency has a heavier tail, and the workflow may need actor, reference, reward, or tool-side models at the same time. + +- Regular vLLM server arguments are passed as `--vllm-*`. For example, vLLM's `--max-model-len` becomes `--vllm-max-model-len`, and `--gpu-memory-utilization` becomes `--vllm-gpu-memory-utilization`. +- Router arguments are passed as `--router-*`. For multi-turn agents that require session affinity, set `--router-policy consistent_hash` so requests for the same `sample.session_id` go to the same worker and improve prefix-cache hit rate; otherwise, vime uses the default `cache_aware` policy. See [Session-Affinity Routing for Multi-Turn Agents](../advanced/vllm-config.md#session-affinity-routing-for-multi-turn-agents). +- Use `--vllm-config` for more complex topologies: PD disaggregation, multi-model serving, heterogeneous server groups, and per-group vLLM overrides. +- For multi-turn or agentic RL, evaluate PD disaggregation. Prefill and decode have different workload shapes, and separating them makes it easier to scale each resource independently. +- For rollout-throughput optimization, also see [Speculative Decoding](../advanced/speculative-decoding.md) and [Low Precision Training](../advanced/low-precision.md). + +## Reference Example + +The full coding-agent example is [`examples/coding_agent_rl`](../_examples_synced/coding_agent_rl/README.md). It shows an end-to-end agent RL setup that is close to a real software-engineering workflow: each sample boots an isolated sandbox, the agent uses tools to edit code, the rollout captures a `git diff`, and a clean sandbox runs the tests to produce the reward. + +This example also demonstrates agent fan-out training. Its middleware splits one trajectory into `subagent`, `wipe` (the chain frozen before compaction), and `final` segments. `generate()` returns `list[Sample]`, and all segments share the same `rollout_id`. + +For a smaller starting point, see [`examples/multi_agent`](../_examples_synced/multi_agent/README.md) for the multi-agent pattern. diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md index 10ace4822..697255531 100644 --- a/docs/en/get_started/customization.md +++ b/docs/en/get_started/customization.md @@ -8,26 +8,26 @@ Below is a summary of all available customization interfaces and their purposes. | Interface Argument | Purpose | | :--- | :--- | -| [`--rollout-function-path`](#1-rollout-function---rollout-function-path) | Override the entire rollout generation logic. | -| [`--custom-generate-function-path`](#2-custom-generate-function---custom-generate-function-path) | Override only the generation step (e.g., for RAG or tool use). | -| [`--custom-rm-path`](#3-reward-model---custom-rm-path) | Implement custom reward computation logic. | -| [`--dynamic-sampling-filter-path`](#4-dynamic-sampling-filter---dynamic-sampling-filter-path) | Filter samples during dynamic sampling (e.g., DAPO). | -| [`--buffer-filter-path`](#5-buffer-filter---buffer-filter-path) | Filter samples in the rollout buffer before training. | -| [`--rollout-sample-filter-path`](#6-rollout-sample-filter---rollout-sample-filter-path) | Determine if individual samples participate in loss calculation. | -| [`--rollout-all-samples-process-path`](#7-rollout-all-samples-process---rollout-all-samples-process-path) | Process all samples (including filtered ones) after rollout. | -| [`--rollout-data-postprocess-path`](#8-rollout-data-postprocess---rollout-data-postprocess-path) | Post-process rollout data after log probs are computed. | -| [`--custom-loss-function-path`](#9-custom-loss-function---custom-loss-function-path) | Implement custom training loss computation. | -| [`--custom-tis-function-path`](#10-custom-tisrs-function---custom-tis-function-path) | Implement custom importance sampling for off-policy correction. | -| [`--custom-pg-loss-reducer-function-path`](#11-custom-pg-loss-reducer---custom-pg-loss-reducer-function-path) | Customize pg_loss reduction (e.g., for Dr.GRPO). | -| [`--custom-reward-post-process-path`](#12-reward-post-processing---custom-reward-post-process-path) | Custom post-processing of rewards before advantage computation. | -| [`--custom-convert-samples-to-train-data-path`](#13-samples-to-train-data-conversion---custom-convert-samples-to-train-data-path) | Override the conversion of samples to training data format. | -| [`--custom-rollout-log-function-path`](#14-logging-functions) | Custom logging for training rollouts. | -| [`--custom-eval-rollout-log-function-path`](#14-logging-functions) | Custom logging for evaluation rollouts. | -| [`--data-source-path`](#15-data-source---data-source-path) | Override the data source for rollout prompts. | -| [`--eval-function-path`](#16-evaluation-function---eval-function-path) | Override the rollout function specifically for evaluation. | -| [`--custom-megatron-init-path`](#17-megatron-hooks) | Custom initialization after Megatron setup. | -| [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hooks) | Custom logic before log probability computation. | -| [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hooks) | Custom logic before each training step. | +| [`--rollout-function-path`](#rollout-function-path) | Override the entire rollout generation logic. | +| [`--custom-generate-function-path`](#custom-generate-function-path) | Override only the generation step (e.g., for RAG or tool use). | +| [`--custom-rm-path`](#custom-rm-path) | Implement custom reward computation logic. | +| [`--dynamic-sampling-filter-path`](#dynamic-sampling-filter-path) | Filter samples during dynamic sampling (e.g., DAPO). | +| [`--buffer-filter-path`](#buffer-filter-path) | Filter samples in the rollout buffer before training. | +| [`--rollout-sample-filter-path`](#rollout-sample-filter-path) | Determine if individual samples participate in loss calculation. | +| [`--rollout-all-samples-process-path`](#rollout-all-samples-process-path) | Process all samples (including filtered ones) after rollout. | +| [`--rollout-data-postprocess-path`](#rollout-data-postprocess-path) | Post-process rollout data after log probs are computed. | +| [`--custom-loss-function-path`](#custom-loss-function-path) | Implement custom training loss computation. | +| [`--custom-tis-function-path`](#custom-tis-function-path) | Implement custom importance sampling for off-policy correction. | +| [`--custom-pg-loss-reducer-function-path`](#custom-pg-loss-reducer-function-path) | Customize pg_loss reduction (e.g., for Dr.GRPO). | +| [`--custom-reward-post-process-path`](#custom-reward-post-process-path) | Custom post-processing of rewards before advantage computation. | +| [`--custom-convert-samples-to-train-data-path`](#custom-convert-samples-to-train-data-path) | Override the conversion of samples to training data format. | +| [`--custom-rollout-log-function-path`](#logging-functions) | Custom logging for training rollouts. | +| [`--custom-eval-rollout-log-function-path`](#logging-functions) | Custom logging for evaluation rollouts. | +| [`--data-source-path`](#data-source-path) | Override the data source for rollout prompts. | +| [`--eval-function-path`](#eval-function-path) | Override the rollout function specifically for evaluation. | +| [`--custom-megatron-init-path`](#megatron-hooks) | Custom initialization after Megatron setup. | +| [`--custom-megatron-before-log-prob-hook-path`](#megatron-hooks) | Custom logic before log probability computation. | +| [`--custom-megatron-before-train-step-hook-path`](#megatron-hooks) | Custom logic before each training step. | ## Agentic workflows through customization interfaces @@ -37,18 +37,18 @@ For most agentic use cases, **start with `--custom-generate-function-path` plus | If you need to … | Use | | :--- | :--- | -| Run a custom agent loop, tool calls, RAG, sandbox execution, browser/terminal interaction, or multi-turn generation for each sample, while reusing vime's default rollout loop | [`--custom-generate-function-path`](#2-custom-generate-function---custom-generate-function-path) | -| Compute verifier rewards, test-based rewards, environment success checks, rule-based rewards, or call an external reward service | [`--custom-rm-path`](#3-reward-model---custom-rm-path) | -| Replace the entire rollout orchestration (only when per-sample customization is not enough) | [`--rollout-function-path`](#1-rollout-function---rollout-function-path) | -| Control task sampling, buffering, requeueing, or custom prompt/task sources | [`--data-source-path`](#15-data-source---data-source-path) | -| Attach custom loss masks, metadata, or convert agentic outputs into training data | [`--rollout-data-postprocess-path`](#8-rollout-data-postprocess---rollout-data-postprocess-path), [`--custom-convert-samples-to-train-data-path`](#13-samples-to-train-data-conversion---custom-convert-samples-to-train-data-path) | -| Debug long-running custom generation, verifier calls, tool calls, or sandbox steps | trace utilities in [`vime.utils.trace_utils`](../developer_guide/trace.md) | +| Run a custom agent loop, tool calls, RAG, sandbox execution, browser/terminal interaction, or multi-turn generation for each sample, while reusing vime's default rollout loop | [`--custom-generate-function-path`](#custom-generate-function-path) | +| Compute verifier rewards, test-based rewards, environment success checks, rule-based rewards, or call an external reward service | [`--custom-rm-path`](#custom-rm-path) | +| Replace the entire rollout orchestration (only when per-sample customization is not enough) | [`--rollout-function-path`](#rollout-function-path) | +| Control task sampling, buffering, requeueing, or custom prompt/task sources | [`--data-source-path`](#data-source-path) | +| Attach custom loss masks, metadata, or convert agentic outputs into training data | [`--rollout-data-postprocess-path`](#rollout-data-postprocess-path), [`--custom-convert-samples-to-train-data-path`](#custom-convert-samples-to-train-data-path) | +| Debug long-running custom generation, verifier calls, tool calls, or sandbox steps | trace utilities in [`vime.observability.trace_utils`](../developer_guide/trace.md) | Native examples of this pattern: [`examples/multi_agent`](../../../examples/multi_agent/README.md) (a `--rollout-function-path`-based multi-agent pattern) and [`examples/fully_async`](../../../examples/fully_async/README.md) (long-tail agentic generation), both keeping vime's default `vllm_rollout` outer loop. ## Detailed Interface Reference -### 1. Rollout Function (`--rollout-function-path`) +### `--rollout-function-path` **Default**: `vime.rollout.vllm_rollout.generate_rollout` @@ -64,11 +64,11 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False) -> Rollout - Adding custom sampling strategies - Integrating external tools or APIs during generation -**Example**: See [examples/multi_agent/rollout_with_multi_agents.py](../../../examples/multi_agent/rollout_with_multi_agents.py) +**Example**: See [examples/fully_async](../_examples_synced/fully_async/README.md) --- -### 2. Custom Generate Function (`--custom-generate-function-path`) +### `--custom-generate-function-path` **Default**: `None` (uses built-in generate function) @@ -76,7 +76,7 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False) -> Rollout **Signature**: ```python -async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample +async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample | list[Sample] ``` **Use Cases**: @@ -84,9 +84,41 @@ async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample - Adding retrieval-augmented generation (RAG) - Multi-turn conversation handling +#### Returning multiple training samples for one prompt + +In agentic settings such as subagents, multi-agent execution, or context compaction, one prompt rollout can naturally split into multiple trainable segments. For example, a subagent trajectory and the main-agent continuation may both need to be trained, or the context before and after compaction may be represented as separate segments. + +You do not need to replace the whole rollout function for this. A `custom_generate` function may return `list[Sample]`. The key contract is that sibling samples produced by the same rollout must share the same `rollout_id`, so vime keeps them together for train-step splitting and loss aggregation instead of counting them as independent rollouts. + +```python +import copy + +from vime.utils.types import Sample + + +async def custom_generate(args, sample: Sample, sampling_params: dict) -> list[Sample]: + segments = await run_agent_and_split_segments(args, sample, sampling_params) + rollout_id = sample.rollout_id if sample.rollout_id is not None else sample.index + + samples: list[Sample] = [] + for segment in segments: + s = copy.copy(sample) + s.tokens = segment.tokens + s.response = segment.response + s.response_length = segment.response_length + s.loss_mask = segment.loss_mask + s.reward = segment.reward + s.status = Sample.Status.COMPLETED + s.rollout_id = rollout_id + samples.append(s) + return samples +``` + +If one full trajectory has a single total reward but is split into `K` training segments, a common pattern is to distribute that reward across the segments, for example by assigning `reward / K` to each segment, so the same rollout reward is not amplified. + --- -### 3. Reward Model (`--custom-rm-path`) +### `--custom-rm-path` **Default**: `None` (uses built-in reward models based on `--rm-type`) @@ -118,7 +150,7 @@ async def batched_custom_rm(args, samples: list[Sample]) -> list[float] --- -### 4. Dynamic Sampling Filter (`--dynamic-sampling-filter-path`) +### `--dynamic-sampling-filter-path` **Default**: `None` @@ -146,7 +178,7 @@ class DynamicFilterOutput: --- -### 5. Buffer Filter (`--buffer-filter-path`) +### `--buffer-filter-path` **Default**: `None` @@ -164,7 +196,7 @@ def buffer_filter(args, rollout_id, buffer: list[list[Sample]], num_samples: int --- -### 6. Rollout Sample Filter (`--rollout-sample-filter-path`) +### `--rollout-sample-filter-path` **Default**: `None` @@ -183,7 +215,7 @@ def filter_function(args, samples: list[Sample]) -> None --- -### 7. Rollout All Samples Process (`--rollout-all-samples-process-path`) +### `--rollout-all-samples-process-path` **Default**: `None` @@ -200,7 +232,7 @@ def process_function(args, samples: list[list[Sample]], data_source) -> None --- -### 8. Rollout Data Postprocess (`--rollout-data-postprocess-path`) +### `--rollout-data-postprocess-path` **Default**: `None` @@ -217,7 +249,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 9. Custom Loss Function (`--custom-loss-function-path`) +### `--custom-loss-function-path` **Default**: `None` (requires `--loss-type custom_loss`) @@ -230,7 +262,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 10. Custom TIS/RS Function (`--custom-tis-function-path`) +### `--custom-tis-function-path` **Default**: `None` @@ -244,7 +276,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 11. Custom pg_loss Reducer (`--custom-pg-loss-reducer-function-path`) +### `--custom-pg-loss-reducer-function-path` **Default**: `None` @@ -266,7 +298,7 @@ def get_pg_loss_reducer( --- -### 12. Reward Post-Processing (`--custom-reward-post-process-path`) +### `--custom-reward-post-process-path` **Default**: `None` (uses default GRPO normalization) @@ -278,7 +310,7 @@ def get_pg_loss_reducer( --- -### 13. Samples to Train Data Conversion (`--custom-convert-samples-to-train-data-path`) +### `--custom-convert-samples-to-train-data-path` **Default**: `None` (uses built-in conversion logic) @@ -318,7 +350,7 @@ dict: { --- -### 14. Logging Functions +### Logging functions #### Training Rollout Logging (`--custom-rollout-log-function-path`) @@ -340,7 +372,7 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool --- -### 15. Data Source (`--data-source-path`) +### `--data-source-path` **Default**: `vime.rollout.data_source.RolloutDataSourceWithBuffer` @@ -369,7 +401,7 @@ class CustomDataSource(DataSource): --- -### 16. Evaluation Function (`--eval-function-path`) +### `--eval-function-path` **Default**: Same as `--rollout-function-path` @@ -381,7 +413,7 @@ class CustomDataSource(DataSource): --- -### 17. Megatron Hooks +### Megatron hooks #### Megatron Initialization (`--custom-megatron-init-path`) @@ -421,6 +453,26 @@ Stabilize MoE RL training by recording and replaying expert routing decisions to | `--use-routing-replay` | Forward-backward routing consistency in training. ([arXiv:2507.18071](https://arxiv.org/abs/2507.18071)) | | `--use-rollout-routing-replay` | R3: Replay routing from rollout during training. Supported by vime's default `vllm_rollout` path. ([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) | +--- + +### 19. Disk Weight-Sync Post-Write Hook (`--custom-update-weight-post-write-path`) + +**Signature**: +```python +def hook(args, version_dir: str, rollout_engines) -> None +``` + +**Purpose**: Called on each trainer rank after a disk weight sync's files are written +(`--update-weight-transport disk`, full or delta mode), before the engines read them. Use it to +publish the writes on a non-POSIX shared filesystem — e.g. upload pending writes to the +backing object store — where another host cannot see the files without an explicit sync. The hook is called +on every rank and must gate itself (e.g. once per container). + +The post-write hook must make the completed version directory visible before it +returns. Host-local full-checkpoint copies then use that published directory as +their source. See [Delta Weight Sync](../advanced/delta-weight-sync.md) for the +delta mechanism. + ## Testing Custom Function Paths vime also provides CPU-only contract tests for customization interfaces. These tests resolve components through import-path strings, so they can validate both built-in hooks and user-defined implementations passed through the same CLI arguments used by training. @@ -446,9 +498,8 @@ python -m pytest \ tests/plugin_contracts/test_plugin_runtime_hook_contracts.py ``` -Each test file can also be executed directly with `python tests/plugin_contracts/.py`, which keeps them compatible with `run-ci-changed`. - -A dedicated `run-ci-plugin-contracts` CI label is also available. Adding it to a PR triggers all four contract test files in parallel (no GPU required). +Each test file can also be executed directly with `python tests/plugin_contracts/.py`. +Buildkite runs all four contract files in its always-on `plugin-contracts` CPU step. For user-defined implementations, you can either export environment variables such as `VIME_CONTRACT_ROLLOUT_FUNCTION_PATH` and `VIME_CONTRACT_CUSTOM_RM_PATH`, or pass overrides directly when running a test file, for example: diff --git a/docs/en/get_started/qa.md b/docs/en/get_started/qa.md index 53d6ab62c..47f03b0f8 100644 --- a/docs/en/get_started/qa.md +++ b/docs/en/get_started/qa.md @@ -65,4 +65,4 @@ 13. **Gradient becomes NaN or Inf during training.** - You can try setting the `--no-check-for-nan-in-loss-and-grad` flag to skip the corresponding training steps. \ No newline at end of file + You can try setting the `--no-check-for-nan-in-loss-and-grad` flag to skip the corresponding training steps. diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md index c9bdac830..134ebcaae 100644 --- a/docs/en/get_started/quick_start.md +++ b/docs/en/get_started/quick_start.md @@ -9,18 +9,26 @@ Since vime may contain temporary patches for vllm/megatron, to avoid potential e ### Hardware Support -**vime** supports multiple NVIDIA GPU hardware platforms: +**vime** supports multiple hardware platforms. -- **B200 Series**: Fully supported with identical setup steps as H-series GPUs +**NVIDIA GPU**: + +Currently stable, production-ready hardware includes: + +- **GB200 / GB300 / B200 / 300 Series**: Fully supported with identical setup steps as H-series GPUs - **H-Series (H100/H200)**: Official support with comprehensive CI testing and stable performance **Important Notes**: - Latest Docker images are compatible with both B-series and H-series GPUs without additional configuration - Megatron backend on H-series GPUs has CI protection, thoroughly validated, recommended for production environments - B-series basic functionality is stable and suitable for development/testing, but currently lacks CI protection -- Both hardware platforms use identical installation and startup procedures +- Both NVIDIA hardware platforms use identical installation and startup procedures +- Other GPUs (e.g., A100/A800) may also run, but are not actively maintained + + +**AMD GPU**: -- For scenarios where Docker is not convenient, please refer to [build_conda.sh](https://github.com/vllm-project/vime/blob/main/build_conda.sh). +See [AMD Usage Tutorial](../platform_support/amd_tutorial.md). ### Pull and Start Docker Container @@ -28,12 +36,12 @@ Please execute the following commands to pull the latest image and start an inte ```shell # Pull the latest image -docker pull inferactinc/public:vime-latest +docker pull vllm/vime:latest # Start the container docker run --rm --gpus all --ipc=host --shm-size=16g \ --ulimit memlock=-1 --ulimit stack=67108864 \ - -it inferactinc/public:vime-latest /bin/bash + -it vllm/vime:latest /bin/bash ``` ### Install vime @@ -53,7 +61,7 @@ You can download required models and datasets from platforms like Hugging Face, ```bash # Download model weights (Qwen3-4B) -hf download zai-org/Qwen3-4B --local-dir /root/Qwen3-4B +hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B # Download training dataset (dapo-math-17k) hf download --repo-type dataset zhuzilin/dapo-math-17k \ @@ -288,8 +296,8 @@ OPTIMIZER_ARGS=( ### VLLM_ARGS: vLLM Service Parameters This part of parameters is used to configure the vLLM inference service. -- `--rollout-num-gpus-per-engine`: Equivalent to vLLM's `tp_size`. -- Other vLLM parameters can be passed to vime by adding the `--vllm-` prefix, and vime will automatically forward them to vLLM. For example, to set vLLM's `--log-level INFO` parameter, just use `--vllm-log-level INFO`. +- `--rollout-num-gpus-per-engine`: Total worker GPUs used by one rollout engine. It equals vLLM's `tensor_parallel_size` only when data and pipeline parallelism are both 1. +- Other vLLM parameters can be passed to vime by adding the `--vllm-` prefix, and vime will automatically forward them to vLLM. For example, to set vLLM's `--uvicorn-log-level info` parameter, use `--vllm-uvicorn-log-level info`. > ⚠️ **Note**: > vime uses `vllm-router` to schedule multiple vLLM engines. `dp_size` is calculated through `rollout-num-gpus / rollout-num-gpus-per-engine`. @@ -304,7 +312,7 @@ VLLM_ARGS=( ### Colocated Actor and Rollout -Under the default configuration, training (Actor) and inference (Rollout) resources are specified separately. Ray allocates `actor_num_nodes * actor_num_gpus_per_node` GPUs to the training part and `rollout_num_gpus` GPUs to inference, that is, training and inference are separated. +Under the default configuration, training (Actor) and inference (Rollout) resources are specified separately. Ray allocates `actor_num_nodes * actor_num_gpus_per_node` GPUs to the training part and `rollout_num_gpus` GPUs to inference, that is, training and inference are separated. When `--rollout-num-gpus` is explicitly set to `0`, vime still parses vLLM arguments and launches the router, but does not launch local vLLM servers. **Standard (Disaggregated) Configuration**: ```bash @@ -318,7 +326,7 @@ ray job submit ... \ In the above configuration, Actor uses 4 cards, and Rollout also uses 4 cards, running in parallel. **Training-Inference Integration (Colocated) Configuration**: -To deploy training and inference on the same group of GPUs, please add the `--colocate` parameter. After enabling, `--rollout-num-gpus` will be ignored to make the number of cards for training and inference equal. +To deploy training and inference on the same group of GPUs, please add the `--colocate` parameter. By default, this makes the number of cards for training and inference equal. You can explicitly set a different positive `--rollout-num-gpus`, for example to use more rollout GPUs than actor GPUs; the extra GPUs are used as rollout-only resources. If `--rollout-num-gpus 0` is set explicitly, vime launches only the router and no local vLLM servers. ```bash ray job submit ... \ @@ -520,7 +528,7 @@ CUSTOM_ARGS=( ## Multi-Node Training for Large-Scale MOE Models -To start a multi-node task, you need to first start a Ray cluster. On node 0, run: +If you use Ray for multi-node training, one option is to start the cluster as follows: ```bash # Node0 (HEAD) @@ -557,3 +565,8 @@ export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=$(ip -o -4 addr show | awk '$4 ~ /^10\. vime has been deeply optimized for distributed training of large-scale Mixture of Experts (MoE) models. We provide some end-to-end training cases for reference: - [Example: Qwen3-30B-A3B with 8xH100](../examples/qwen3-30B-A3B.md) +- [Example: 8xH100 Training GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md) +- [Example: 32xH100 Training GLM-5.2](../examples/glm5.2-744B-A40B.md) +- [Example: 64xH100 Training GLM-4.7](../examples/glm4.7-355B-A32B.md) +- [Example: 128xH100 Training DeepSeek-R1](../examples/deepseek-r1.md) +- Scripts such as `scripts/run_qwen3_30b_a3b.py` and `scripts/run_glm45_355b_a32b.py` also support multi-node training. Their documentation is still being expanded. diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md index 84fafe43f..9d73e9564 100644 --- a/docs/en/get_started/usage.md +++ b/docs/en/get_started/usage.md @@ -18,22 +18,21 @@ There are four main parameters for cluster resource allocation: - `--actor-num-nodes`: The number of nodes required for RL actor training. - `--actor-num-gpus-per-node`: The number of GPUs per node for RL actor training. - - `--rollout-num-gpus`: The total number of GPUs required for rollout (inference). - - `--rollout-num-gpus-per-engine`: The number of GPUs per inference engine. This parameter is similar to vLLM's `tp_size`. When performing multi-node serving, this value should be the total number of GPUs. For example, if serving one model with 2 nodes and 16 GPUs, this value should be 16. + - `--rollout-num-gpus`: The total number of GPUs required for rollout (inference). Set it to `0` to still parse vLLM arguments and launch the router without launching local vLLM servers. + - `--rollout-num-gpus-per-engine`: The total worker GPU count for one inference engine. It equals vLLM's `tensor_parallel_size` only when data and pipeline parallelism are both 1. For example, if one model is served across 2 nodes and 16 GPUs, this value should be 16. With the default configuration, we use these parameters to allocate `actor_num_nodes * actor_num_gpus_per_node` GPUs for training and `rollout_num_gpus` GPUs for inference via Ray, thus achieving a separation of training and inference resources. For co-located training and inference, you also need to configure: - - `--colocate`: Enables co-located training and inference. When enabled, it ignores `--rollout-num-gpus` and makes the number of GPUs for training and inference equal. + - `--colocate`: Enables co-located training and inference. By default, this makes the number of GPUs for training and inference equal. You can explicitly set a different positive `--rollout-num-gpus`, for example to use more rollout GPUs than actor GPUs; the extra GPUs are used as rollout-only resources. If `--rollout-num-gpus 0` is set explicitly, vime launches only the router and no local vLLM servers. Additionally, vime supports Prefill and Decode disaggregation (PD Disaggregation). You can set the number of servers used for Prefill by setting the `--prefill-num-servers` argument. ### Choosing Training Backend -vime supports multiple training backends, which can be selected via the `--train-backend` parameter: - -- `megatron` (default): Uses Megatron-LM as the training backend, supporting efficient training of large-scale models. +vime currently uses Megatron-LM as its training backend. The compatibility option +`--train-backend megatron` may still be supplied explicitly. ### Loading Megatron @@ -145,13 +144,14 @@ Note: - Before the first training step, vime will synchronize the parameters from Megatron to vLLM. Therefore, the `--hf-checkpoint` does not need to contain the latest training parameters, and you do not need to change the HF checkpoint when resuming training. - By default, vLLM reads the maximum context length from the `config.json` in the Hugging Face checkpoint. You can use the `--vllm-max-model-len` parameter to override this value to support longer inference. - During co-located training and inference, although Megatron and vLLM will offload sequentially, they still need to leave some memory for each other. You need to adjust vLLM's total VRAM usage by reducing `--vllm-gpu-memory-utilization`. - - vime supports passing through vllm-router parameters by adding a `router` prefix to the original parameter name. For example, vllm-router's `--balance-abs-threshold` parameter should be set as `--router-balance-abs-threshold`. Since vllm-router uses cache-aware routing by default, it may cause uneven request distribution. You can set `--router-balance-abs-threshold 0` to force balanced distribution, but this may affect prefix cache hit rate in multi-turn conversation scenarios. + - vime supports passing through vllm-router parameters by adding a `router` prefix to the original parameter name. For example, vllm-router's `--balance-abs-threshold` parameter should be set as `--router-balance-abs-threshold`. Since vllm-router uses cache-aware routing by default, it may cause uneven request distribution. You can set `--router-balance-abs-threshold 0` to force balanced distribution, but this may affect prefix cache hit rate in multi-turn conversation scenarios. For multi-turn sessions that require session affinity, set `--router-policy consistent_hash` and send a stable `x-session-id` for each session. + - If vLLM engines are pre-launched by an external system, connect to them with `--rollout-external-engine-addrs host1:port host2:port`. When the trainer and engines cannot form an NCCL weight-update group, use `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`; vime writes a complete HF checkpoint and asks vLLM to hot-load it through `update_weights_from_disk`. For large models or cross-cluster deployments, use `--update-weight-mode delta --update-weight-transport disk` instead. See [External Rollout Engines Roadmap](../advanced/external-rollout-engines.md) and [Delta Weight Sync](../advanced/delta-weight-sync.md). For details on some of vLLM's customizations and the principles behind how vime incorporates vLLM, please see the "How to Use vLLM" section. ### Data Format -Currently, vime only supports loading files in `.jsonl` format, where each line of the file is a JSON object. An example of a single data entry (expanded) is as follows: +vime supports `.jsonl` and `.parquet` files; reading Parquet requires `pyarrow`. Each record in either format should contain the fields selected by `--input-key` and `--label-key`. An expanded JSONL record looks like this: ```json { @@ -177,11 +177,26 @@ This corresponds to the following configuration: Please note that the `step_loss_mask` (default=1) here is for SFT phase. If it is set to 0, the turn will not contibute to the final loss; if it is set to 1, vime will use the normal `loss_mask`. Additionally, we provide a `metadata_key`, which defaults to `"metadata"`. When read, vime will load the metadata from the data, which can be helpful for custom data generation or creating custom reward models. +If one run mixes multiple data sources, put `source_name` in the sample metadata: + +```json +{ + "prompt": "...", + "label": "...", + "metadata": { + "source_name": "math" + } +} +``` + +The recommended contract is to put the source identifier in `metadata["source_name"]`; vime also recognizes a dynamically set `sample.source` from custom data sources. When rollout samples are converted to training data, vime carries one `source_names` entry per sample to the training side. The source lookup order is dynamic `sample.source`, then `metadata["source_name"]`; if neither is set, the source is `"unknown"`. This is useful for custom rewards, filters, logging, and future per-source routing such as OPD teacher selection. + ### Hyperparameters for RL Training - `--advantage-estimator`: Specifies the RL algorithm for the training process. Currently supported algorithms include: - `grpo` ([https://arxiv.org/abs/2402.03300](https://arxiv.org/abs/2402.03300)) - `gspo` ([https://arxiv.org/abs/2507.18071](https://arxiv.org/abs/2507.18071)) + - `cispo` ([https://arxiv.org/abs/2506.13585](https://arxiv.org/abs/2506.13585)) - `reinforce_plus_plus` and `reinforce_plus_plus_baseline` ([https://arxiv.org/abs/2501.03262](https://arxiv.org/abs/2501.03262)) - `ppo` ([https://arxiv.org/abs/1707.06347](https://arxiv.org/abs/1707.06347)) - `--calculate-per-token-loss`: By default, vime calculates loss on a per-sample basis, i.e., `mean(sum(sample_i) / len(sample_i))`. Enable this flag to calculate loss on a per-token basis, i.e., `sum(sum(sample_i)) / sum(len(sample_i))`. @@ -219,35 +234,17 @@ To use PPO, set: --advantage-estimator ppo ``` -**Note: In PPO, the Critic and Actor request GPUs in parallel**, which should be considered when allocating resources. Specifically: - -- The critic model occupies a separate set of GPUs, independent from the actor's GPU resources. -- You can configure critic resources using `--critic-num-nodes` and `--critic-num-gpus-per-node`. -- If critic resource parameters are not configured, the same resource configuration as the actor will be used by default. +**Note: In PPO, the critic and actor share the same training GPU group.** You do not need to reserve a separate set of GPUs for the critic. Specifically: -Cluster resource allocation example: - -```bash -# Actor uses 1 node, 4 GPUs ---actor-num-nodes 1 ---actor-num-gpus-per-node 4 - -# Critic uses 1 node, 4 GPUs (parallel to Actor) ---critic-num-nodes 1 ---critic-num-gpus-per-node 4 - -# Rollout uses 8 GPUs ---rollout-num-gpus 8 -``` +- PPO creates separate actor and critic training process groups, but places them on the same train placement group. +- The critic training scale follows the actor configuration, and the actor / critic Megatron parallel topology must currently stay identical. +- PPO forces train-side offload so that actor and critic can wake up and release memory on the same GPUs in turn. +- There are currently no separate CLI arguments for configuring critic training resources; the critic node count and GPUs per node are derived from the actor configuration. -With the above configuration, a total of `4 (actor) + 4 (critic) + 8 (rollout) = 16` GPUs are required. PPO-related parameters: -- `--critic-load`: Checkpoint path for the critic model. -- `--critic-save`: Save path for the critic model. -- `--critic-lr`: Learning rate for the critic model. -- `--critic-lr-warmup-iters`: Number of warmup steps for the critic model. +- `--megatron-config-path`: YAML config for role-specific Megatron overrides, such as setting critic-specific `load`, `save`, `lr`, or warmup parameters. - `--num-critic-only-steps`: Number of steps to train only the critic at the beginning of training. - `--eps-clip`: PPO clip range. - `--value-clip`: Clip range for value loss. @@ -326,7 +323,6 @@ vime supports customizing data generation (rollout) to various degrees. output = await post( f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate", { - "model": args.hf_checkpoint, "token_ids": prompt_token_ids, "sampling_params": {"max_tokens": sampling_params["max_new_tokens"]}, } @@ -409,7 +405,7 @@ Each model gets its own router. The per-model router info is accessible via `arg **Server group features:** - `worker_type`: `regular`, `prefill`, `decode`, or `placeholder` (reserves GPU slots without creating engines) - `overrides`: Dict of vLLM `EngineArgs` field overrides applied on top of `--vllm-*` CLI args -- `num_gpus_per_engine`: Per-group TP size override +- `num_gpus_per_engine`: Per-group total worker GPU count override ## How to Use Megatron diff --git a/docs/en/index.rst b/docs/en/index.rst index 324d31fcf..c8c552acf 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -12,6 +12,21 @@ vime is built on `slime `_, the RL framework beh - DeepSeek V3 series (DeepSeek V3, V3.1, DeepSeek R1); - Llama 3. +Start by Use Case +----------------- + +- New to vime: :doc:`get_started/quick_start` +- Configure training and rollout arguments: :doc:`get_started/usage` +- Add custom generation, reward, or rollout functions: :doc:`get_started/customization` +- Build agentic RL workflows: :doc:`get_started/agent` +- Configure production vLLM rollout topology: :doc:`advanced/vllm-config` +- Connect external rollout engines: :doc:`advanced/external-rollout-engines` +- Sync weights as byte-level deltas: :doc:`advanced/delta-weight-sync` +- Use PD disaggregation: :doc:`advanced/pd-disaggregation` +- Use BF16 training with FP8 rollout or FP8 KV cache: :doc:`advanced/low-precision` +- Understand CI and reliability coverage: :doc:`developer_guide/ci` +- Debug, trace, and profile long-running jobs: :doc:`developer_guide/debug`, :doc:`developer_guide/trace`, :doc:`developer_guide/profiling` + .. toctree:: :maxdepth: 1 :caption: Get Started @@ -26,21 +41,29 @@ vime is built on `slime `_, the RL framework beh :caption: Dense examples/qwen3-4B.md + examples/glm4-9B.md .. toctree:: :maxdepth: 1 :caption: MoE examples/qwen3-30B-A3B.md + examples/glm5.2-744B-A40B.md + examples/glm4.7-355B-A32B.md + examples/deepseek-r1.md .. toctree:: :maxdepth: 1 :caption: Advanced Features + advanced/on-policy-distillation.md advanced/speculative-decoding.md advanced/reproducibility.md advanced/fault-tolerance.md + advanced/observability.md advanced/pd-disaggregation.md + advanced/external-rollout-engines.md + advanced/delta-weight-sync.md advanced/vllm-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md @@ -51,6 +74,7 @@ vime is built on `slime `_, the RL framework beh _examples_synced/fully_async/README.md _examples_synced/multi_agent/README.md + _examples_synced/coding_agent_rl/README.md .. toctree:: :maxdepth: 1 @@ -60,3 +84,9 @@ vime is built on `slime `_, the RL framework beh developer_guide/debug.md developer_guide/trace.md developer_guide/profiling.md + +.. toctree:: + :maxdepth: 1 + :caption: Hardware Platforms + + platform_support/amd_tutorial.md diff --git a/docs/en/platform_support/amd_tutorial.md b/docs/en/platform_support/amd_tutorial.md new file mode 100644 index 000000000..5ddc1a0d6 --- /dev/null +++ b/docs/en/platform_support/amd_tutorial.md @@ -0,0 +1,82 @@ +# Quick Start + +This document will guide you through setting up the environment and getting started with vime on AMD ROCm, covering environment configuration, data preparation, weight conversion, and training startup. + +## Basic Environment Setup + +### Pull and Start Docker Container + +Execute the following commands to pull the latest ROCm image and start a persistent container: + +```shell +# Pull the ROCm image +docker pull vllm/vime-rocm + +# Start the container +docker run -d --name vime --ulimit nofile=1048576:1048576 \ + --ipc=host --network=host --device=/dev/kfd --device=/dev/dri \ + --security-opt seccomp=unconfined --group-add video --privileged \ + -e WANDB_API_KEY=$wandb_key vllm/vime-rocm +# wandb key is optional if you want to track with WandB + +# Enter the container +docker exec -it vime bash +``` + +## Model and Dataset Download + +Download the required model and training dataset using `huggingface_hub`: + +```bash +# Download model weights (Qwen3-8B) +hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B + +# Download training dataset (dapo-math-17k) +hf download zhuzilin/dapo-math-17k --repo-type dataset --local-dir /root/dapo-math-17k +``` + +## Model Weight Conversion + +### HF → Megatron torch_dist ckpt + +Use Vime's built-in Hugging Face-to-Megatron loader for conversion. Load the model configuration for Qwen3-8B, then run the conversion. Two ROCm-specific flags are required: `--no-gradient-accumulation-fusion` and `--attention-backend flash`. + +```bash +cd /root/vime && source scripts/models/qwen3-8B.sh + +HIP_VISIBLE_DEVICES=0 PYTHONPATH=/root/vime:/root/Megatron-LM \ + torchrun --nproc-per-node=1 tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \ + --no-gradient-accumulation-fusion --attention-backend flash \ + --hf-checkpoint /root/Qwen3-8B --save /root/Qwen3-8B_torch_dist +``` + +> **Note**: On ROCm, use `HIP_VISIBLE_DEVICES` in place of `CUDA_VISIBLE_DEVICES` to select GPUs. The `--attention-backend flash` and `--no-gradient-accumulation-fusion` flags are required to avoid issues during conversion. + +## Training + +Run training using the ROCm-specific script. `VISIBLE_GPUS` specifies which GPUs to use, and `NUM_ROLLOUT` controls the total number of sampling→training rounds. The script automatically unsets NVTE environment variables that are not needed on ROCm. + +```bash +NUM_ROLLOUT=100 VISIBLE_GPUS=0,1 bash scripts/run-qwen3-8B-amd.sh +``` + +### Configuration Notes + +- **VISIBLE_GPUS** + - Two free GPU indices. + - Script masks execution to these GPUs, avoids clashes. + - Uses: + - TP = 2 + - Single vLLM engine + - Colocate mode + - DP = 1 +- **NUM_ROLLOUT** + - Number of training steps. + - Default is **3** for a smoke test. +- Each run requires approximately **230 GB** across the two selected GPUs. +- Launch only on GPUs with sufficient free memory. + +> **Final note**: After finishing the run, if rerunning with a different `NUM_ROLLOUT`, make sure to clear the save directory to avoid mismatch error. +```bash +rm -rf /root/Qwen3-8B_vime/ +``` diff --git a/docs/zh/advanced/arch-support-beyond-megatron.md b/docs/zh/advanced/arch-support-beyond-megatron.md index 82a32779b..f7e93b0a7 100644 --- a/docs/zh/advanced/arch-support-beyond-megatron.md +++ b/docs/zh/advanced/arch-support-beyond-megatron.md @@ -22,8 +22,8 @@ Megatron 的模型实例化分为两步：首先根据配置生成“层规格 * **对应文件**: `vime_plugins/models/hf_attention.py` 3. **对齐模型权重** - 模型结构跑通后，还需要确保权重能正确加载。我们借助 [mbridge](https://github.com/ISEEKYAN/mbridge) 库，通过 `Qwen3NextBridge` 建立了 HuggingFace Checkpoint 与 Megatron 参数之间的命名映射关系，实现双向互通。 - * **对应文件**: `vime_plugins/mbridge/qwen3_next.py` + 模型结构跑通后，还需要确保权重能正确加载。vime 将 HuggingFace 到 Megatron 的名称映射和 tensor 变换直接放在 checkpoint loader 旁。 + * **对应文件**: `vime/backends/megatron_utils/hf_to_megatron/qwen3_next.py` 通过这三层协同，我们成功地将一个 Megatron 原本不支持的复杂模型结构（以其 HuggingFace 实现为载体），运行在了 Megatron 的并行框架之上，并完整保留了模型并行、MoE 加速、流水线调度等全部关键能力。 diff --git a/docs/zh/advanced/delta-weight-sync.md b/docs/zh/advanced/delta-weight-sync.md new file mode 100644 index 000000000..8a1a33d3f --- /dev/null +++ b/docs/zh/advanced/delta-weight-sync.md @@ -0,0 +1,59 @@ +# Delta 权重同步 + +Delta 权重同步只发送两次同步之间发生变化的字节，而不是每次都写一份完整 checkpoint，以此让非 colocate 的 rollout engine 保持最新。它面向大模型、跨集群或跨数据中心的训推解耦场景——这种场景下每次都写整份 actor 权重是主要开销。 + +它**只支持 disk transport**。训练端把每次同步发布为一份 canonical HF checkpoint 目录；engine 的 `/pull_weights` 端点（随 vime 的 vllm patch 提供）把 apply 扇出到 **engine 覆盖的每一个 host** 并校验，随后 engine 通过**原生**的 `update_weights_from_disk` 端点 reload 打过补丁的本地 checkpoint。vime 对每个 engine 只与一个端点通信，所以多节点 serving 和外部 rollout engine 在 vime 侧都不需要任何额外支持。 + +## 配置 + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/delta-updates +--update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt +--update-weight-delta-encoding xor # 或: overwrite +--update-weight-delta-checksum xxh3-128 # 或: blake3, adler32 +``` + +| 参数 | 作用 | +|---|---| +| `--update-weight-disk-dir` | 训练端发布 delta、rollout host 读取 delta 的共享文件系统目录。 | +| `--update-weight-local-checkpoint-dir` | host 本地（如 NVMe）的完整 HF checkpoint，由 `/pull_weights` 保持同步——delta 原地 apply，发布的完整 checkpoint 则整份替换。每个 host 在第一次 `/pull_weights` 时由 engine 的 model path seed。 | +| `--update-weight-delta-encoding` | 磁盘上的 delta 编码：`xor`（默认）或 `overwrite`。 | +| `--update-weight-delta-checksum` | 逐 tensor 完整性 checksum：`xxh3-128`（默认）、`blake3` 或 `adler32`。 | + +delta 始终用 zstd（level 1）压缩；profiling 显示对这类数据它在 wire 大小和解压速度上都优于 lz4 / gzip / snappy / brotli，所以不做成可配置项。 + +## 工作原理 + +1. **Seed。** 第一次同步时，训练端为每个参数捕获一份 CPU snapshot——从 `--hf-checkpoint` seed，而这正是每个 rollout host 物化本地 checkpoint 的来源。此次不发布任何东西；这份 snapshot 就是下一次同步 diff 的基准。训练端同时发出 `target_version=0` 的 `/pull_weights`，让每个 host 现在就物化本地 base，与 snapshot 捕获重叠进行。 +2. **Publish。** 之后每次同步，训练端把每个 gather 出的 HF tensor 与 snapshot 做 diff，编码、压缩，写到 `--update-weight-disk-dir` 下的新版本目录 `weight_v{N:06d}/`。该目录是一份 canonical HF checkpoint——`model-NNNNN.safetensors` 文件装着压缩后的 diff tensor，外加 `model.safetensors.index.json`（tensor 名 → 文件）承载 apply 元数据——所以这个产物是可移植的，不绑定训练端的并行 layout。随后 snapshot 推进到新值，供下次 diff。 +3. **Pull。** 训练端对每个 engine 调用 `/pull_weights`。engine 内部把请求广播到每个节点的每个 rank；每个 host 把新版本的 delta 原地 apply 进它的本地 checkpoint（host 级文件锁把同 host 的多个 rank 合并成一次 apply）。apply 在 tensor 之间并行，并逐 tensor 校验（见"完整性"）；只有**每一个 host** 都持有校验通过的 checkpoint，该调用才报告成功。 + + `/pull_weights` 并不绑定 delta：每个发布的版本是自描述的。若某个版本是一份普通的完整 HF + checkpoint（index 中没有 delta 元数据），pull 就直接整份拷贝——同时重置链条，因此晚加入的 + 新 host 从最近的完整版本 seed，而不必回放全部 delta，旧的 delta 也可以被清理。vime 的 + full 模式 disk 同步在设置了 `--update-weight-local-checkpoint-dir` 时正是走这条路径。 +4. **Reload。** engine 通过原生 `update_weights_from_disk` 路径 reload 打过补丁的本地 checkpoint——权重加载代码从不接触 delta 格式。 + +由于 snapshot 是从 `--hf-checkpoint`（engine 真正的 base）seed，而不是从当前 GPU 权重 seed，即使 Megatron→HF 往返不是逐字节相等（例如 embedding / LM head 中被裁掉的 vocab padding 行），该方案对任意模型也都正确。 + +## 编码 + +两种编码都是字节级、与 dtype 无关的，所以量化 checkpoint 也走同一条路径。engine 从每个版本的 index 元数据读取所用编码。 + +- **`xor`**（默认）：写 `new ^ old`。wire 最小、apply 最快（顺序访问、对 cache 友好；未变化的字节是 0，被压缩器压到极小）。它是一个对合（involution），所以必须**恰好对正确的 base apply 一次**——apply 两次会还原。 +- **`overwrite`**：写变化的位置及其新的绝对值。wire 更大、apply 是对 cache 不友好的分散写，但**幂等**：重复 apply（或把部分 apply 的 delta 补完）无论执行多少次都收敛到同一状态。当“可重复 apply”比 wire 大小更重要时用它。 + +## 完整性 + +训练端把每个 tensor 新状态的逐 tensor checksum 存进版本里。apply 之后每个 host 重新计算 checksum，**任何不匹配都会 raise**——失败会通过 `/pull_weights` 的响应传回，所以损坏的 delta 或错误的 base 会直接报错失败，而不会把坏权重提供出去。apply 还拒绝乱序执行：一个版本只会在它声明的 base 版本之上 apply。 + +`--update-weight-delta-checksum` 选择算法。checksum 不是 apply 的瓶颈（apply 受解压 + XOR 限制），所以这是一个 digest 属性的选择，而非速度选择：`xxh3-128`（默认）是最宽的快速非加密 digest；`blake3` 是加密 digest，用于不可信存储；`adler32` 用于与期望它的系统互操作。 + +## 共享文件系统可见性 hook + +在 POSIX 共享文件系统（NFS、Lustre……）上不需要额外步骤。对于需要显式 commit/refresh 才能让写入跨 host 可见的对象存储挂载，可以提供两个可选 hook（通过 import 路径加载——vime 和 vllm 里都不存在任何厂商特定代码）： + +- `--custom-update-weight-post-write-path`（vime，训练端）：在一个版本的文件写完之后、通知 engine 读取之前调用（例如把待写入数据上传到底层对象存储）。签名：`hook(args, version_dir, rollout_engines)`。 +- `--custom-update-weight-pre-read-path`（vime，engine 端）：在每个 host 上、`/pull_weights` 读取 delta 目录之前于 engine 内部调用（例如刷新挂载视图）。签名：`hook(delta_dir, target_version)`。 diff --git a/docs/zh/advanced/external-rollout-engines.md b/docs/zh/advanced/external-rollout-engines.md new file mode 100644 index 000000000..49dfa85b1 --- /dev/null +++ b/docs/zh/advanced/external-rollout-engines.md @@ -0,0 +1,110 @@ +# External Rollout Engines 配置路线图 + +External rollout engine 指的是：vLLM engine 不由 vime 训练任务启动，而是由外部系统预先部署和管理；vime 只在训练时连接这些 engine，注册 router，并在需要时同步训练后的 actor 权重。 + +这篇文档是一个导航页。它帮助你判断什么时候该用 `--rollout-external-engine-addrs`，什么时候该继续使用 `--vllm-config`，以及 external 场景下该选择 full checkpoint update from disk 还是 delta update。 + +## 从哪里开始 + +| 目标 | 推荐入口 | +| :--- | :--- | +| engine 已经由外部系统启动，只想让 vime 连接并做 rollout | `--rollout-external-engine-addrs` | +| engine 仍由 vime 启动，但需要 PD 分离、多模型、异构 server group 或 per-group overrides | [vLLM Config](vllm-config.md) | +| 训练器和 external engine 可以建立 NCCL group | 默认的 `--update-weight-mode full --update-weight-transport nccl` | +| 训练器和 external engine 不能建立 NCCL group，但能共享同一路径的文件系统 | `--update-weight-mode full --update-weight-transport disk` | +| 大模型跨集群或跨数据中心同步，full checkpoint 太重 | `--update-weight-mode delta --update-weight-transport disk` | +| rollout serving 想使用独立 vLLM 环境，甚至不同型号或不同厂家的 GPU | external engine + disk transport | +| 需要 reference、reward、tool-side model 等冻结模型 | 优先用 [vLLM Config](vllm-config.md#3-多模型服务) 的 `update_weights: false` | + +delta mode 仅支持 disk transport。通过 NCCL 同步权重时请使用 full mode。 + +## External Engine 做了什么 + +使用 external engine 时，先独立启动 vLLM server： + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10090 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10091 ... +``` + +训练任务里传入这些地址： + +```bash +python train.py \ + --rollout-external-engine-addrs host1:10090 host2:10091 \ + ... +``` + +vime 会请求每个 engine 的 `/server_info` 或 `/get_server_info`，推断 engine 的 GPU 数、TP/PP 信息和 worker 类型（`regular`、`prefill`、`decode`）。如果没有提供 `--vllm-router-ip/--vllm-router-port`，vime 会启动自己的 router，并把这些 external engine 注册进去。 + +这条路径适合 serving 生命周期由训练任务外部管理的部署：例如独立的推理集群、跨 Ray 集群部署、手工预热的 vLLM engine，或由其他编排系统管理的 rollout service。 + +## 与 `--vllm-config` 的关系 + +`--rollout-external-engine-addrs` 和 `--vllm-config` 互斥，因为它们负责不同的边界： + +- `--vllm-config`：vime 负责 engine 生命周期。你用 YAML 描述 topology，vime 启动 server group、router，并管理多模型和选择性权重更新。 +- `--rollout-external-engine-addrs`：外部系统负责 engine 生命周期。vime 只发现已启动的 engine，接入 router，并把它们当作默认 rollout model。 + +如果你的主要需求是多模型 serving、reference/reward 冻结模型、PD 分离或异构组配置，优先使用 `--vllm-config`。如果 engine 已经在训练任务外部部署好，再使用 external engine。 + +## 环境与硬件解耦 + +External engine 的一个重要含义是：vLLM serving 侧不需要使用 vime 训练任务的 Python 环境、Megatron 环境或 Ray runtime。它可以运行在单独的 vLLM 容器、独立集群或其他编排系统里；vime 只依赖 HTTP endpoint、`/server_info` 信息，以及所选权重同步方式需要的通信路径。 + +当使用 disk transport 时，权重通过共享文件系统上的 HF checkpoint 或 safetensors delta 传递，再由 vLLM 通过 `update_weights_from_disk` 热加载。这条路径不要求训练 GPU 和 rollout GPU 是同一型号，甚至不要求是同一厂家；只要 vLLM 本身支持该硬件后端、模型格式和精度配置即可。例如训练可以在一组 GPU 上运行，rollout serving 可以放在另一组不同型号或不同厂家的 GPU 上。 + +如果使用 NCCL transport，则仍然需要满足 NCCL 通信和硬件兼容性要求。跨厂家、跨不兼容网络或跨数据中心部署通常应选择 `--update-weight-transport disk`。 + +## Update From Disk + +full checkpoint update from disk 是 external 场景最简单的兜底路径： + +```bash +--update-weight-mode full +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/full-updates +``` + +每次权重同步时，训练端会在 `--update-weight-disk-dir` 下写一个完整 HF checkpoint 目录，例如 `weight_v000123/`，然后通过 HTTP 调用每个 vLLM engine 的 `update_weights_from_disk`，让 engine 在不重启进程的情况下重新加载 checkpoint。 + +额外设置 `--update-weight-local-checkpoint-dir` 后，每个 engine 会先把发布的 checkpoint pull 到它覆盖的每个 host 的本地磁盘（`/pull_weights`，随 vime 的 vllm patch 提供），再从本地（如 NVMe）reload——共享文件系统每个 host 只读一次，而不是每个 rank 读一次；当共享目录是对象存储或 engine 跨多个节点时尤其重要。 + +这个模式的优点是控制面简单：不要求训练器和 engine 建 NCCL group，只要求二者能看到同一个共享文件系统路径。缺点也直接：每次同步都写完整 actor 权重，对大模型和高频同步来说非常重。 + +调试时可以加： + +```bash +--update-weight-disk-keep-files +``` + +这样 vime 不会在 engine 确认加载后清理完整 checkpoint 目录，方便检查写出的 HF checkpoint。 + +## Update With Delta + +delta update 面向大模型、跨集群或跨数据中心训推解耦。它不每次都写完整 checkpoint，而是在训练端保留上一次同步的 CPU snapshot，逐参数比对，只发布变化的字节；每个 engine 的 `/pull_weights` 端点（随 vime 的 vllm patch 提供）把 delta apply 进 engine 覆盖的每个 host 的本地 checkpoint，再通过原生 `update_weights_from_disk` 端点 reload。vime 只调用 engine 的 HTTP 端点，所以多节点 external engine 与 vime 拉起的 engine 行为一致。 + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-disk-dir /shared/fs/delta-updates +--update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt +``` + +机制、编码、完整性校验以及共享文件系统可见性 hook 详见 [Delta 权重同步](delta-weight-sync.md)。 + +## 部署检查清单 + +- external engine 的 HTTP 地址必须能从训练任务访问。 +- external engine 可以使用独立 vLLM 环境；不需要安装 vime 或 Megatron 训练环境。 +- disk transport 支持训练和 rollout 使用不同型号或不同厂家的 GPU，前提是 vLLM 支持对应硬件和模型格式。 +- disk transport 要求训练端和 vLLM engine 看到同一个 `--update-weight-disk-dir` 路径；路径只在训练端可见是不够的。 +- external engine 当前不支持 vime 的 fault tolerance 恢复流程；engine 生命周期由外部系统负责。 +- `--vllm-config` 与 `--rollout-external-engine-addrs` 互斥。 +- delta mode 不支持 `--colocate`，因为 colocated 权重同步通过 CUDA IPC 传句柄，delta 编码不会节省实际传输量。 + +## 参考工作 + +[Cursor Research Team 的 Composer 2 技术报告](https://arxiv.org/html/2603.24477v2) 公开描述了一个相近的生产形态：训练和 rollout generation 高度异步，Cursor 与 Fireworks AI 合作运行 RL inference；每个训练 step 都把更新后的权重写到共享 S3，并用 delta compression 降低传输量，不同区域的 inference 集群再从共享 delta chain 下载并重建权重。 + +vime 的 external engine、update from disk 和 delta disk transport 面向同一类基础设施问题：训练与推理解耦后，权重同步必须能跨进程、跨集群甚至跨数据中心工作，同时不能让训练主循环被完整模型传输拖住。 diff --git a/docs/zh/advanced/fault-tolerance.md b/docs/zh/advanced/fault-tolerance.md index d230f9d4d..b22f322cd 100644 --- a/docs/zh/advanced/fault-tolerance.md +++ b/docs/zh/advanced/fault-tolerance.md @@ -12,8 +12,8 @@ vime 当前提供 rollout-engine fault tolerance： -- 对 vLLM rollout engine 做 health check； -- heartbeat timeout 后重启 rollout engine； +- 对 vLLM rollout server 做 health check； +- heartbeat timeout 后重启 rollout server； - 重启后正确更新参数； - 保存 debug rollout dump，用于不重新跑 rollout 的情况下 replay 训练侧问题； - trace/profiling hook，用于检查 long-tail rollout 行为。 @@ -22,7 +22,7 @@ vime 当前提供 rollout-engine fault tolerance： ## Rollout Health Checks -rollout 过程中，vime 会定期向所有 vLLM engine 发送 heartbeat 请求（`/health`）。如果 heartbeat timeout，异常 vLLM engine 会被停止。当前 rollout round 完成后，vime 会重启 engine，并在其继续服务后续 rollout 请求前更新到正确参数。 +rollout 过程中，vime 会定期向所有 vLLM server 发送 heartbeat 请求（`/health`）。如果 heartbeat timeout，异常 vLLM server 会被停止。当前 rollout round 完成后，vime 会重启 server，并在其继续服务后续 rollout 请求前更新到正确参数。 主要参数： @@ -65,7 +65,7 @@ rollout 过程中，vime 会定期向所有 vLLM engine 发送 heartbeat 请求 - 如果大 MoE 模型启动阶段 health check 失败，增大 `--rollout-health-check-first-wait`。 - 如果短暂负载高峰导致误判，增大 `--rollout-health-check-timeout`。 -- 如果某个 engine 在 weight sync 后反复重启，检查 vLLM log 和最近的 rollout debug dump。 +- 如果某个 server 在 weight sync 后反复重启，检查 vLLM log 和最近的 rollout debug dump。 - 如果失败发生在 trainer 而非 rollout，从 checkpoint 恢复，并用 debug replay 确认保存的 rollout batch 是否有效。 ## 相关文档 diff --git a/docs/zh/advanced/low-precision.md b/docs/zh/advanced/low-precision.md new file mode 100644 index 000000000..bf18621c1 --- /dev/null +++ b/docs/zh/advanced/low-precision.md @@ -0,0 +1,143 @@ +# 低精度训练与 Rollout + +vime 中的低精度能力主要用于让 rollout 更快、更省显存，同时保持训练侧数值稳定。对于大规模 MoE RL，推荐的生产路径是： + +> **Megatron BF16 训练 + vLLM FP8 rollout/inference** + +Megatron 侧保持 BF16/torch_dist 的可训练 checkpoint；vLLM 侧使用 FP8 Hugging Face checkpoint 做 rollout。权重更新时，vime 会根据 `--hf-checkpoint` 中的量化配置，把更新后的 BF16 权重量化后再发送给 vLLM。 + +## Feature Maturity + +| 功能 | 状态 | 推荐用法 | +|---|---|---| +| BF16 training + FP8 rollout/inference | Stable | 大规模 MoE RL 的默认推荐路径。训练保持稳定，rollout 降低显存和带宽开销。 | +| vLLM rollout FP8 KV cache | Stable，取决于当前 vLLM 版本和 GPU stack 支持 | 通过 `--vllm-kv-cache-dtype fp8_e4m3` 提升 long-context 或 agentic rollout 的 KV cache 容量。 | +| INT4 rollout / INT4 QAT | Beta | 当 rollout 显存或吞吐压力很高，并且目标模型路径已经验证时使用。 | +| FP8 training + FP8 rollout | Experimental | 适合研究训推不一致和吞吐优化，但仍有 optimizer/checkpoint 相关限制。 | + +## BF16 训练 + FP8 Rollout + +这是 vime 当前最主要的生产路径。 + +你可以通过将 `--hf-checkpoint` 指向 blockwise quantized Hugging Face checkpoint 来开启 FP8 rollout。可以用如下命令从 BF16 checkpoint 转换： + +```bash +python tools/convert_hf_to_fp8.py \ + --model-dir $BF16_MODEL \ + --save-dir $FP8_MODEL \ + --strategy block --block-size 128 128 \ + --max-workers 4 +``` + +请确保转换后的 checkpoint 中 `config.json` 包含正确的 `quantization_config`。vime 会在权重更新时使用这个配置，因此训练侧可以保持 BF16，而 rollout 侧收到 FP8 权重。 + +示例： + +```bash +# Megatron 训练 checkpoint 仍然是 BF16 / torch_dist。 +--ref-load /path/to/model_torch_dist + +# vLLM rollout checkpoint 使用 FP8 Hugging Face。 +--hf-checkpoint /path/to/model-fp8-hf +``` + +## Rollout 使用 FP8 KV Cache + +对于 long-context、multi-turn 或 agentic workload，KV cache 容量经常是瓶颈。由于 vime 通过 `--vllm-` 前缀透传 vLLM 参数，可以直接开启 FP8 KV cache： + +```bash +--vllm-kv-cache-dtype fp8_e4m3 +``` + +这是 rollout 侧配置，不会改变 Megatron 的训练精度。它可以提升 vLLM 的有效 KV cache 容量，从而支持更长 context 或更高并发；实际精度和性能表现取决于当前 vLLM 版本与 GPU stack。 + +## FP8 训练 + FP8 Rollout + +vime 也支持 experimental 的 FP8 training 路径。我们观察到，在一些设置下同时使用 FP8 training 和 FP8 inference，可以提升推理吞吐并降低训推不一致。更多细节请参考 [这篇博客](https://lmsys.org/blog/2025-11-25-fp8-rl/)。 + +### 快速开始 + +1. 使用 `tools/convert_hf_to_fp8.py` 将 Hugging Face 模型权重转换为 FP8 格式。 + +2. 添加 FP8 训练参数： + +```bash +--fp8-format e4m3 +--fp8-recipe blockwise +# --fp8-param-gather # 可选；目前与 CPU Adam 不兼容 +``` + +3. 确保设置 `NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1`。vime 默认会为 Ray actors 设置为 `1`。 + +4. 启动 FP8 training example： + +```bash +# Qwen3-4B FP8 training +bash scripts/low_precision/run-qwen3-4b-fp8.sh + +# Qwen3-30B-A3B FP8 training (2 nodes) +bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh +``` + +### 实现说明 + +1. 如果启用 FP8 recipe，TransformerEngine layers 会在 FP8 context 中构建。 +2. 训练时，权重和 activation 会在线量化为 NVFP8 格式，并在 forward/backward GEMM 中使用 cuBLAS FP8 GEMM。 +3. RL 权重更新时，Megatron 会先把 FP8 权重反量化为 BF16，然后 vime 再把 BF16 权重量化为 FP8 并发送给 vLLM。 +4. 从 training engine 保存 checkpoint 时，会反量化回 BF16 并保存为 `torch_dist`。 + +目前只有 TransformerEngine 中的 `Linear` 和 `GroupLinear` 层使用 FP8。`embedding` 和 `lm_head` 保持原始精度。如果未开启 `--fp8-param-gather`，TransformerEngine 中的权重以 BF16 存储，仅在 `GEMM` 或 `GroupGEMM` 时转换为 FP8。 + +### 已知限制 + +`--fp8-param-gather` 可以节省显存，但目前需要 TransformerEngine `FusedAdam`，这与大规模 Megatron-LM RL 中常用的 CPU Adam offload 路径冲突。 + +## INT4 QAT 训练 + +INT4 STE（Straight-Through Estimator）训练和 INT4 inference 可以进一步降低 rollout 显存并提升吞吐。在目标模型和 reward setup 验证前，请把这条路径视作 beta。 + +### 快速开始 + +1. 将 Hugging Face 权重转换为 INT4： + +```bash +python tools/convert_hf_to_int4_direct.py \ + --model-dir /path/to/your/original/models \ + --save-dir /path/to/your/save/models +``` + +如果只需要 INT4 rollout，把 `--hf-checkpoint` 指向转换后的 INT4 checkpoint 即可。 + +2. 开启 INT4 fake QAT： + +```json +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" +``` + +`OPEN_TRAINING_INT4_GROUP_SIZE` 通常设置为： + +- `128`：`moonlight-16B-A3B`、`qwen3-30B-A3B`、`qwen3-235B-A22B-int4`； +- `32`：`kimi-k2-Thinking-int4`。 + +3. 启动 example： + +```bash +# Moonlight-16B-A3B INT4 training +bash scripts/low_precision/run-moonlight-16B-A3B-int4.sh + +# Qwen3-30B-A3B INT4 training +bash scripts/low_precision/run-qwen3-30B-A3B-int4.sh + +# Qwen3-235B-A22B INT4 training (8 nodes) +bash scripts/low_precision/run-qwen3-235B-A22B-int4.sh + +# Kimi-k2-Thinking INT4 training (32 nodes) +bash scripts/low_precision/run-kimi-k2-Thinking-int4.sh +``` + +多机环境请根据集群配置启动 Ray 服务。 diff --git a/docs/zh/advanced/megatron-config.md b/docs/zh/advanced/megatron-config.md index 224592e8e..484d353e4 100644 --- a/docs/zh/advanced/megatron-config.md +++ b/docs/zh/advanced/megatron-config.md @@ -74,7 +74,6 @@ megatron: ```bash python train.py \ --advantage-estimator ppo \ - --use-critic \ --megatron-config-path megatron_ppo.yaml \ --tensor-model-parallel-size 2 \ --sequence-parallel \ @@ -84,14 +83,13 @@ python train.py \ --expert-tensor-parallel-size 1 \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 8 \ - --critic-num-nodes 1 \ - --critic-num-gpus-per-node 8 \ ... ``` 在这个模式下： -- CLI 负责共享的并行策略和资源配置； +- `--advantage-estimator ppo` 会自动启用 critic，不需要额外的 `--use-critic` 参数； +- CLI 负责共享的并行策略和资源配置；当前 PPO 下 critic 的训练资源会跟随 actor 配置； - YAML 负责 actor / critic 的差异项，比如 `lr`、`load`、`save`、optimizer 或 scheduler 相关参数。 ### 只覆盖一个角色 @@ -114,6 +112,7 @@ megatron: - **目前只支持 PPO。** `--megatron-config-path` 当前主要用于 PPO 工作流中的 actor / critic 角色配置。对于 GRPO、REINFORCE++ 等不依赖 critic 的流程，目前不建议使用这套角色配置。 - **当前 PPO 下，actor 和 critic 的 Megatron 并行配置必须一致。** 特别是 `tensor_model_parallel_size`、`pipeline_model_parallel_size`、`context_parallel_size`、`expert_model_parallel_size`、`expert_tensor_parallel_size`、`sequence_parallel` 等拓扑相关参数，不应在 actor 和 critic 之间配置成不同的值。 +- **当前 PPO 下，actor 和 critic 共享同一组 train placement group。** critic 的节点数和每节点 GPU 数由 actor 配置派生，不能作为独立资源规模配置。 - **推荐把并行相关参数继续放在 CLI 中。** 当前最稳妥的用法是：并行与资源参数写在公共 CLI 中，只在 YAML 中覆盖角色差异项，例如 `lr`、`load`、`save`、warmup、optimizer / scheduler 参数等。 如果你在 actor 和 critic 之间写入不同的并行拓扑，当前行为不受支持，可能导致初始化或训练过程出错。 @@ -126,6 +125,6 @@ megatron: 可以。缺失角色会自动继承公共 CLI 参数，不需要把所有参数都重复写一遍。 -### Q: 可以把 `--actor-num-nodes` 或 `--critic-num-gpus-per-node` 写进 YAML 吗？ +### Q: 可以把资源配置写进 YAML 吗？ -不可以。当前资源分配和 placement group 仍由 CLI 参数控制，YAML 中对应字段会被忽略。 \ No newline at end of file +不可以。当前资源分配和 placement group 仍由 CLI 参数控制，YAML 中对应字段会被忽略。其中 `--actor-num-nodes` / `--actor-num-gpus-per-node` 决定 PPO 的 train 资源规模；critic 的节点数和每节点 GPU 数会跟随 actor 配置，不能独立配置。 diff --git a/docs/zh/advanced/observability.md b/docs/zh/advanced/observability.md new file mode 100644 index 000000000..4108a6b3a --- /dev/null +++ b/docs/zh/advanced/observability.md @@ -0,0 +1,97 @@ +# 观测 + +vime 的默认观测路径很简单：训练指标继续进 W&B / TensorBoard；vLLM 的高频 Prometheus metrics 不再上传 W&B；request timing 从 vLLM response `meta_info` 写进 sample trace，并在 rollout 结束时聚合成少量 `perf/...` 指标。 + +## W&B / TensorBoard 里会看到什么 + +W&B 和 TensorBoard 仍然记录 reward、loss、KL、entropy、eval 等训练指标。额外的 vLLM request timing 会放在 `perf/` 下，例如： + +```text +perf/request/e2e_latency/mean +perf/request/queue_time/median +perf/request/count +perf/request/profiled_count +perf/decode/throughput/mean +perf/prefill/bootstrap_queue_duration/mean +perf/prefill/bootstrap_duration/mean +perf/prefill/alloc_wait_duration/mean +perf/prefill/forward_duration/max +perf/prefill/transfer_speed_gb_s/mean +perf/decode/prealloc_duration/mean +perf/decode/bootstrap_duration/mean +perf/decode/alloc_wait_duration/mean +perf/decode/transfer_duration/max +perf/decode/forward_duration/mean +``` + +这些指标是每个 rollout step 聚合一次，不是每个 request 上报一次，所以不会像上传完整 Prometheus metrics 那样拖慢 W&B。 + +不开 PD 时仍然会有通用的 `perf/request/...` 和可用的 `perf/decode/throughput/...`。`perf/prefill/...` 和更细的 `perf/decode/...duration` 只有在 vLLM 返回对应 `pd_*` timing 字段时才会出现。 + +## Prometheus metrics 存在哪里 + +vime 自己不存 Prometheus 的每秒数据。vLLM / router 只暴露 `/metrics` 和 `/engine_metrics` HTTP endpoint；Prometheus 定期 scrape 这些 endpoint，并把时间序列写进 Prometheus 自己的 TSDB。 + +因此： + +- 如果没有启动 Prometheus，这些 serving metrics 只存在于 vLLM 进程内存和当前 endpoint 输出里，不会形成历史记录。 +- 如果启动了 Prometheus，历史数据存放在 Prometheus 的 `--storage.tsdb.path`。 +- vime 不把这些高频 metrics 上传到 W&B。 + +常用的 vLLM metrics 包括： + +```text +vllm:num_queue_reqs +vllm:num_running_reqs +vllm:num_prefill_bootstrap_queue_reqs +vllm:num_prefill_inflight_queue_reqs +vllm:num_decode_prealloc_queue_reqs +vllm:num_decode_transfer_queue_reqs +vllm:kv_transfer_speed_gb_s_bucket +vllm:kv_transfer_latency_ms_bucket +vllm:kv_transfer_total_mb_bucket +``` + +这些适合在 Grafana / Prometheus 里看实时 queue buildup、transfer speed、latency histogram、失败计数等 serving 状态。 + +## 如何启动 Prometheus + +Prometheus 需要在训练运行时启动，因为它只能 scrape 当前正在暴露的 endpoint，不能在训练结束后从 vLLM endpoint 里补回过去的数据。它不需要放进训练 Python 进程里，推荐作为同一台机器或同一个作业里的旁路进程运行。 + +最小配置如下，把 `ROUTER_IP:ROUTER_PORT` 换成 vime 日志里的 router 地址，或者用户显式设置的 `--vllm-router-ip` / `--vllm-router-port`： + +```yaml +global: + scrape_interval: 10s + +scrape_configs: + - job_name: vime-vllm + metrics_path: /engine_metrics + static_configs: + - targets: + - "ROUTER_IP:ROUTER_PORT" +``` + +启动 Prometheus 时把 TSDB 目录放到持久化路径上： + +```bash +prometheus \ + --config.file=/path/to/prometheus.yml \ + --storage.tsdb.path=/path/to/prometheus-data \ + --storage.tsdb.retention.time=7d \ + --web.listen-address=0.0.0.0:9090 +``` + +vime 镜像里已经安装了 `prometheus`，可以直接在容器里用上面的命令启动。也可以从同一个镜像再起一个旁路容器，只要它能访问 router 地址，并把 `/path/to/prometheus-data` 挂到持久化目录即可。 + +如果 `--storage.tsdb.path` 在容器本地盘里，容器回收后数据也会丢；如果挂到 NFS、持久化卷或作业输出目录，训练结束后可以重新启动 Prometheus 指向同一个 TSDB 目录，再用 Prometheus UI 或 Grafana 查询历史时间段。这里的“回放”是时间序列回放和图表分析，不是 per-request trace 的完整重放；per-sample request timing 仍然走 sample trace / debug rollout 数据。 + +## Trace viewer + +`--save-debug-rollout-data` 保存的 sample trace 会包含 vLLM `meta_info` 里的 timing 字段。trace viewer 直接读取这些 attrs，并用 `pd_*` 字段展示 `[P]` / `[D]` 虚拟 lane。 + +```bash +python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt +``` + +默认路径不需要单独保存 `ReqTimeStats(...)` 日志，也不需要 Loki 或 compact 工具。 diff --git a/docs/zh/advanced/on-policy-distillation.md b/docs/zh/advanced/on-policy-distillation.md new file mode 100644 index 000000000..9ee810770 --- /dev/null +++ b/docs/zh/advanced/on-policy-distillation.md @@ -0,0 +1,128 @@ +# 在策略蒸馏 (On-Policy Distillation) + +在策略蒸馏 (OPD) 使用学生当前策略采样的 response token 来训练学生。对于学生轨迹中访问到的每个前缀，固定的教师模型为同一个 next token 评分，从而沿学生自己的轨迹提供稠密的 token 级学习信号。在 vime 中，这一信号以逆 KL 的采样估计惩罚 advantage，因此可以与 GRPO、PPO、REINFORCE++ 等 advantage estimator 组合；当任务奖励为零时，同一机制就是纯蒸馏。 + +## 关键参数 + +| 参数 | 说明 | +|------|------| +| `--use-opd` | 启用在策略蒸馏。使用 OPD 的必需标志。 | +| `--opd-type` | OPD 类型：`vllm` 或 `megatron`。启用 `--use-opd` 时必须设置。 | +| `--opd-kl-coef` | OPD KL 惩罚系数（默认值：1.0）。控制蒸馏信号相对于 RL advantage 的权重。 | +| `--opd-teacher-load` | 教师模型的 Megatron checkpoint 路径。`--opd-type=megatron` 时**必须**设置，`--opd-type=vllm` 时**不可**设置。 | +| `--opd-teacher-ckpt-step` | 可选的教师模型 checkpoint 步数。 | +| `--opd-teacher-model` | `--opd-type=vllm` 时发送给外部 VLLM 教师服务的可选模型名。 | + +## 原理 + +记 $\pi_\theta$ 为学生策略，$\pi_T$ 为教师策略，$h_t$ 为学生生成轨迹中采样 token $a_t$ 之前的历史。按照 [Thinking Machines Lab 给出的定义](https://thinkingmachines.ai/blog/on-policy-distillation/)，token 级逆 KL 为 + +$$ +D_{\mathrm{KL}}\left(\pi_\theta(\cdot \mid h_t) \| \pi_T(\cdot \mid h_t)\right) += \mathbb{E}_{a_t \sim \pi_\theta(\cdot \mid h_t)}\left[ +\log \pi_\theta(a_t \mid h_t) - \log \pi_T(a_t \mid h_t) +\right]. +$$ + +这里的顺序很重要：KL 的第一个参数是学生分布，期望同样对学生分布取值。教师不生成训练轨迹，而是评估学生实际采样的 token。 + +vime 不会遍历完整词表来精确计算这个期望。对于每个采样 token，它使用如下 Monte Carlo 贡献： + +$$ +\hat d_t = \log \pi_\theta(a_t \mid h_t) - \log \pi_T(a_t \mid h_t), +\qquad a_t \sim \pi_\theta(\cdot \mid h_t), +$$ + +然后修改基础 advantage： + +$$ +\hat A_t = A_t - \lambda_{\mathrm{opd}}\hat d_t. +$$ + +其中 $A_t$ 来自所配置的 estimator（纯蒸馏时为零），$\lambda_{\mathrm{opd}}$ 是 `--opd-kl-coef`。尽管 KL 的期望非负，单个样本的 $\hat d_t$ 仍可能为负。策略损失使用修改后的 $\hat A_t$，因此 OPD 项与 GRPO、PPO、REINFORCE++、GSPO 等 advantage estimator 的选择相互独立。 + +## 两种教师模式 + +### VLLM 模式 (`--opd-type vllm`) + +教师模型运行在外部 VLLM 服务器上，教师的 log-probs 在 rollout 阶段获取。 + +**适用场景**：教师与学生架构不同，或教师模型太大无法与训练模型同时加载。由于教师需要为学生的原始 token ID 评分，两者仍须使用兼容的 tokenizer 和词表。 + +**工作流程**： +1. 外部 VLLM 服务器运行教师模型。 +2. 在 rollout 阶段，自定义 reward 函数（`vime.rollout.on_policy_distillation.reward_func`）将学生采样的 token ID 发送给教师服务器，并获取教师对这些相同 token 的 log-probability。 +3. 自定义后处理函数（`vime.rollout.on_policy_distillation.post_process_rewards`）将教师 log-probs 裁剪到 response 范围并存储到 `sample.teacher_log_probs` 中。 +4. 在训练阶段，vime 从基础 advantage 中减去按 `--opd-kl-coef` 缩放后的采样 log-probability 差值。 + +**配置**： +```bash +--use-opd +--opd-type vllm +--opd-kl-coef 1.0 +--custom-rm-path vime.rollout.on_policy_distillation.reward_func +--custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards +--rm-url http://:/inference/v1/generate +``` + +### Megatron 模式 (`--opd-type megatron`) + +教师模型通过 `--opd-teacher-load` 直接加载到 Megatron 中，教师的 log-probs 在训练前向传播阶段计算。 + +**适用场景**：教师与学生/参考模型架构相同，且能放入 GPU 显存。 + +**工作流程**： +1. 教师模型在初始化时作为额外的 Megatron 模型加载。 +2. 在训练前向传播阶段，教师模型为每个样本计算 log-probs。 +3. 内联计算 KL 惩罚并应用到 advantages。 + +**配置**： +```bash +--use-opd +--opd-type megatron +--opd-kl-coef 1.0 +--opd-teacher-load /path/to/teacher_torch_dist +``` + +> **注意**：教师 checkpoint 必须是 Megatron 格式（`torch_dist` 或 `torch`）。可以使用 `tools/convert_hf_to_torch_dist.py` 从 HuggingFace 格式转换。 + +## 运行示例 + +完整的示例脚本在 `examples/on_policy_distillation/` 中： + +### VLLM 教师 + +```bash +# 1. 下载模型和数据 +hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B +hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k + +# 2. 转换学生模型 +cd /root/vime +source scripts/models/qwen3-8B.sh +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-8B \ + --save /root/Qwen3-8B_torch_dist + +# 3. 运行 +bash examples/on_policy_distillation/run-qwen3-8B-opd.sh +``` + +### Megatron 教师 + +```bash +# 1. 将学生和教师模型都转换为 Megatron 格式 +# 2. 运行 +bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh +``` + +## 初步结果 + +使用 Qwen3-8B-Base 模型在 [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) 数据集的一部分上进行 SFT，然后在剩余数据上用 Qwen3-32B 教师进行在策略蒸馏，Math500 评测结果如下： + +| | Pass@1 | +|-----------------------------------------------|--------| +| Qwen3-8B-Base + SFT | 76% | +| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% | diff --git a/docs/zh/advanced/pd-disaggregation.md b/docs/zh/advanced/pd-disaggregation.md index fd88f7d1c..7f51f4321 100644 --- a/docs/zh/advanced/pd-disaggregation.md +++ b/docs/zh/advanced/pd-disaggregation.md @@ -30,7 +30,7 @@ vime 支持两种 PD 配置方式。 ### 高级路径：`--vllm-config` -生产级 rollout topology 推荐使用 [vLLM Config](vllm-config.md)。它可以独立配置 prefill 和 decode group，也能表达 EPD-style layout、heterogeneous engine group、multi-model serving 和 per-group vLLM override。 +生产级 rollout topology 推荐使用 [vLLM Config](vllm-config.md)。它可以独立配置 prefill 和 decode group，也能表达 EPD-style layout、heterogeneous server group、multi-model serving 和 per-group vLLM override。 示例： @@ -43,12 +43,13 @@ vllm: num_gpus: 4 num_gpus_per_engine: 2 overrides: - chunked_prefill_size: 8192 + enable_chunked_prefill: true + max_num_batched_tokens: 8192 - worker_type: decode num_gpus: 12 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.88 + gpu_memory_utilization: 0.88 ``` 启动： diff --git a/docs/zh/advanced/reproducibility.md b/docs/zh/advanced/reproducibility.md index 2f0976508..0be85890b 100644 --- a/docs/zh/advanced/reproducibility.md +++ b/docs/zh/advanced/reproducibility.md @@ -50,3 +50,30 @@ bash scripts/run-qwen2.5-0.5B-reproducibility.sh ``` 这个 PR 中记录了 wandb 的截图 [pull#370](https://github.com/THUDM/slime/pull/370)。 + +## Train/rollout log-prob alignment（GLM-5） + +除单侧 bitwise 复现外，vime 还可以对齐训练与 rollout（推理）的 log-prob。目前该能力只支持 **GLM-5 结构**（MLA + DSA sparse attention），并要求 deterministic VLLM、batch-invariant DeepGEMM 与 DeepEP 构建。所需 Megatron 侧对齐 hook 由 Vime 在运行时安装，不需要额外 Megatron patch。 + +Supported in this path: + +- DSA sparse attention (`flashmla_sparse` prefill/decode), including deterministic NSA RadixCache/prefix cache; +- DeepGEMM batch-invariant block-FP8 forward for dense and grouped-MoE layers (with BF16 backward); +- fp32 MoE router (the LM head stays bf16 on both train and rollout — matching precision, not fp32, is what aligns); +- VLLM rollout 使用 DeepEP low-latency，Megatron 训练使用 DeepEP normal。 + 第二次小 payload normal dispatch 保留每个 top-k route，token owner 按 + slot 顺序做 FP32 加权归约；这条对齐路径不支持普通 Megatron all-to-all； +- 支持 bf16 或 FP8-E4M3 KV cache。`flashmla_sparse` 路径把 KV 以 FP8 + packed 格式保存，只 gather 并反量化被选中的 page，再交给 BF16 sparse + kernel。维护的 gate 默认使用 FP8-E4M3，不使用 rollout routing replay + (R3)，因此包括 router 和 experts 在内的主模型参数都会执行 backward； + 辅助 DSA indexer 通过 `--freeze-indexer` 始终保持冻结。 + +回归 gate 是 `tests/test_glm52_6layer_deterministic_e2e.py`（6-layer GLM-5.2， +单机 EP8）：它执行真实的 Megatron→VLLM online-weight-update rollout， +训练全部主模型参数，并断言 `train_rollout_logprob_abs_diff < 1e-6`（已验证的 +DeepEP 对齐参考结果为 `x e-7` 量级）。 + +另有一个较短的 EP8 gate `tests/test_glm52_layerwise_zero_e2e.py`，会同时 +记录训推两侧 decoder layer 0–5 的可见输出，并要求所有匹配 hidden-state +元素的绝对误差严格等于 0。 diff --git a/docs/zh/advanced/speculative-decoding.md b/docs/zh/advanced/speculative-decoding.md index 9e3cb9569..ccd74c8dc 100644 --- a/docs/zh/advanced/speculative-decoding.md +++ b/docs/zh/advanced/speculative-decoding.md @@ -8,14 +8,14 @@ vLLM 把投机采样的所有配置收敛到一个 JSON（`SpeculativeConfig`） `--vllm-speculative-config` 透传。对于有 MTP 层的模型（例如 GLM-4.7、DeepSeek-V3/R1），传入： ```bash ---vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' +--vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' ``` 如果要使用单独训练的 draft model，在同一个 JSON 里加上 `model`（可选还可加 `draft_tensor_parallel_size` 等）： ```bash ---vllm-speculative-config '{"method":"eagle","num_speculative_tokens":3,"model":"/your/draft/model/path"}' +--vllm-speculative-config '{"method":"eagle","num_speculative_tokens":4,"model":"/your/draft/model/path"}' ``` 要从头训练一个 draft model，可以参考 [TorchSpec](https://github.com/lightseekorg/TorchSpec) diff --git a/docs/zh/advanced/vllm-config.md b/docs/zh/advanced/vllm-config.md index 2e6e2c9e3..4aa318472 100644 --- a/docs/zh/advanced/vllm-config.md +++ b/docs/zh/advanced/vllm-config.md @@ -31,11 +31,11 @@ vllm: - name: # 必填。模型的唯一标识符。 model_path: # 可选。HF checkpoint 路径。默认使用 --hf-checkpoint。 update_weights: # 可选。是否从训练同步权重。自动推断。 - num_gpus_per_engine: # 可选。该模型所有组的默认 TP 大小。 + num_gpus_per_engine: # 可选。该模型所有组中单引擎的默认 worker GPU 总数。 server_groups: # 必填。服务器组配置列表。 - worker_type: # 必填。可选：regular、prefill、decode、placeholder。 num_gpus: # 必填。分配给该组的 GPU 总数。 - num_gpus_per_engine: # 可选。该组的 TP 大小覆盖。 + num_gpus_per_engine: # 可选。该组中单引擎的 worker GPU 总数覆盖。 overrides: # 可选。vLLM EngineArgs 字段覆盖。 ``` @@ -48,7 +48,7 @@ vllm: | `name` | `str` | **必填** | 模型唯一名称（如 `"actor"`、`"ref"`、`"reward"`）。用作 `args.vllm_model_routers` 的 key。 | | `model_path` | `str` | `args.hf_checkpoint` | HuggingFace checkpoint 路径。同一模型内的所有服务器组必须使用相同的 model path。 | | `update_weights` | `bool` | 自动推断 | 该模型是否接收训练权重更新。未设置时自动推断：如果 `model_path` 与 `--hf-checkpoint` 匹配则为 `true`，否则为 `false`。 | -| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | 该模型服务器组的默认 TP 大小。各组可单独覆盖。 | +| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | 该模型中单引擎的默认 worker GPU 总数。各组可单独覆盖。 | | `server_groups` | `list` | **必填** | `ServerGroupConfig` 条目列表，定义引擎拓扑。（`engine_groups` 作为向后兼容别名仍可使用。） | #### 服务器组级字段 @@ -57,8 +57,8 @@ vllm: |------|------|--------|------| | `worker_type` | `str` | **必填** | 引擎类型：`regular`（标准）、`prefill`（PD prefill worker）、`decode`（PD decode worker）或 `placeholder`（占位，不启动引擎）。 | | `num_gpus` | `int` | **必填** | 该组的 GPU 总数。必须 > 0。 | -| `num_gpus_per_engine` | `int` | 模型的 `num_gpus_per_engine` | TP 大小覆盖。每个引擎实例的 GPU 数量。 | -| `overrides` | `dict` | `{}` | vLLM `ServerArgs` 字段覆盖。优先级最高，覆盖 `--vllm-*` CLI 参数和模型级默认值。 | +| `num_gpus_per_engine` | `int` | 模型的 `num_gpus_per_engine` | 单个引擎实例的 worker GPU 总数。只有 DP 和 PP 都为 1 时才等于 TP。 | +| `overrides` | `dict` | `{}` | vLLM `EngineArgs` 字段覆盖。优先级最高，覆盖 `--vllm-*` CLI 参数和模型级默认值。 | ### Worker 类型 @@ -107,10 +107,10 @@ vllm: server_groups: - worker_type: prefill num_gpus: 4 - num_gpus_per_engine: 2 # 2 个 prefill 引擎，TP=2 + num_gpus_per_engine: 2 # 默认 DP/PP 下为 2 个 prefill 引擎，TP=2 - worker_type: decode num_gpus: 12 - num_gpus_per_engine: 4 # 3 个 decode 引擎，TP=4 + num_gpus_per_engine: 4 # 默认 DP/PP 下为 3 个 decode 引擎，TP=4 ``` ```bash @@ -125,7 +125,7 @@ python train.py \ - 为 decode 使用更大的 TP（降低延迟） - 独立扩展 prefill 和 decode 的容量 -> **注意：** PD 分离使用 vllm-router (vllm-router)，并设置 `pd_disaggregation=True`。 +> **注意：** PD 分离使用 vllm-router，并设置 `pd_disaggregation=True`。 ### 3. 多模型服务 @@ -174,18 +174,25 @@ from vime.rollout.vllm_rollout import get_model_url from vime.utils.http_utils import post async def my_generate(args, sample, sampling_params): - # 路由到 actor 模型（默认） - actor_url = get_model_url(args, "actor", "/generate") - output = await post(actor_url, {"text": sample.prompt, "sampling_params": sampling_params}) - + # 路由到 actor 模型（默认端点为 /inference/v1/generate） + actor_url = get_model_url(args, "actor") + output = await post(actor_url, { + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, + }) + # output["choices"][0] 含 token_ids、logprobs.content[i].logprob 及 finish_reason + # 路由到 reference 模型 - ref_url = get_model_url(args, "ref", "/generate") - ref_output = await post(ref_url, {"text": sample.prompt, "sampling_params": sampling_params}) - + ref_url = get_model_url(args, "ref") + ref_output = await post(ref_url, { + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, + }) + # 路由到 reward 模型（如 OpenAI 兼容 API） reward_url = get_model_url(args, "reward", "/v1/chat/completions") reward_output = await post(reward_url, {...}) - + ... ``` @@ -232,9 +239,9 @@ vllm: num_gpus: 2 # 预留 2 个 GPU（不创建引擎） ``` -### 6. 按组覆盖 ServerArgs +### 6. 按组覆盖 EngineArgs -使用 `overrides` 将 vLLM `ServerArgs` 字段应用到特定服务器组，而不影响其他组： +使用 `overrides` 将 vLLM `EngineArgs` 字段应用到特定服务器组，而不影响其他组： ```yaml vllm: @@ -244,10 +251,12 @@ vllm: num_gpus: 8 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.85 - context_length: 32768 - chunked_prefill_size: 4096 - enable_torch_compile: true + gpu_memory_utilization: 0.85 + max_model_len: 32768 + enable_chunked_prefill: true + max_num_batched_tokens: 4096 + compilation_config: + mode: 3 ``` 覆盖具有**最高优先级**，会覆盖基础的 `--vllm-*` CLI 参数和模型级默认值。这对以下场景特别有用： @@ -257,7 +266,7 @@ vllm: ### 7. 独立 vLLM 启动器 -虽然 `--vllm-config` 是为 vime 的训练流水线设计的，但它也可以作为纯推理场景的强大启动器，通过 `--rollout-external` 模式或配置 vime 仅关注推理服务。 +虽然 `--vllm-config` 是为 vime 的训练流水线设计的，但它也可以作为纯推理场景的强大启动器，通过外部 engine 地址或配置 vime 仅关注推理服务。 **使用预启动的外部引擎：** @@ -265,17 +274,23 @@ vllm: ```bash # 步骤 1：外部启动 vLLM 引擎 -vllm serve /path/to/model --port 10090 ... -vllm serve /path/to/model --port 10091 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10090 ... +VLLM_SERVER_DEV_MODE=1 vllm serve /path/to/model --port 10091 ... # 步骤 2：将 vime 连接到外部引擎 python train.py \ - --rollout-external \ --rollout-external-engine-addrs host1:10090 host2:10091 \ ... ``` -> **注意：** `--vllm-config` 和 `--rollout-external` 互斥。当你希望 vime 管理完整的引擎生命周期时，使用 `--vllm-config`；当引擎已预部署时，使用 `--rollout-external`。 +vime 会请求每个外部引擎的 `/server_info`，自动推断 +`rollout_num_gpus`、单个 engine 的 GPU 数、vLLM 并行参数，以及 +prefill/decode worker 类型。如果没有提供 `--vllm-router-ip/--vllm-router-port`， +vime 会自己启动 router，并把这些外部引擎注册进去。 + +> **注意：** `--vllm-config` 和 `--rollout-external-engine-addrs` 互斥。当你希望 vime 管理完整的引擎生命周期时，使用 `--vllm-config`；当引擎已预部署时，使用 `--rollout-external-engine-addrs`。 + +关于 external engine 的选择、update from disk 和 delta disk transport，见 [External Rollout Engines 配置路线图](external-rollout-engines.md)。 --- @@ -332,7 +347,7 @@ vime 自动为每个 sample 分配一个唯一的 `session_id`（存储在 `samp | 选项 | 冲突原因 | |------|----------| | `--prefill-num-servers` | PD 分离通过 YAML 中的 `server_groups` 配置 | -| `--rollout-external` | 外部引擎有自己的拓扑；config 在内部管理生命周期 | +| `--rollout-external-engine-addrs` | 外部引擎有自己的拓扑；config 在内部管理生命周期 | --- @@ -351,12 +366,13 @@ vllm: num_gpus: 4 num_gpus_per_engine: 2 overrides: - chunked_prefill_size: 8192 + enable_chunked_prefill: true + max_num_batched_tokens: 8192 - worker_type: decode num_gpus: 12 num_gpus_per_engine: 4 overrides: - mem_fraction_static: 0.88 + gpu_memory_utilization: 0.88 - name: ref model_path: /data/models/Qwen3-32B @@ -392,33 +408,42 @@ python train.py \ **自定义 rollout 函数 (`my_agent/rollout.py`)：** ```python +from transformers import AutoTokenizer + from vime.rollout.vllm_rollout import get_model_url from vime.utils.http_utils import post async def generate_with_models(args, sample, sampling_params): """使用 actor 生成，用 reward 模型打分，与 reference 比较。""" - # 从 actor 生成 - actor_url = get_model_url(args, "actor", "/generate") + # 从 actor 生成（默认端点为 /inference/v1/generate） + actor_url = get_model_url(args, "actor") actor_output = await post(actor_url, { - "text": sample.prompt, - "sampling_params": sampling_params, - "return_logprob": True, + "token_ids": sample.tokens, + "sampling_params": {"max_tokens": 1024, "temperature": 1.0, "top_p": 1.0, "logprobs": 1}, }) - - # 获取 reference logprobs 用于 KL penalty - ref_url = get_model_url(args, "ref", "/generate") + response_ids = actor_output["choices"][0]["token_ids"] + + # 获取 reference logprobs 用于 KL penalty。max_tokens=1 + prompt_logprobs 对提交的 + # token_ids 打分；从顶层 "prompt_logprobs" 字段读取。 + ref_url = get_model_url(args, "ref") ref_output = await post(ref_url, { - "text": sample.prompt + actor_output["text"], - "sampling_params": {"max_new_tokens": 0, "temperature": 0}, - "return_logprob": True, + "token_ids": sample.tokens + response_ids, + "sampling_params": {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 1}, }) - - # 用 reward 模型打分 + + # 用 reward 模型为 actor response 打分（OpenAI 兼容） + tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + response_text = tokenizer.decode(response_ids, skip_special_tokens=True) + prompt_messages = ( + sample.prompt + if isinstance(sample.prompt, list) + else [{"role": "user", "content": sample.prompt}] + ) reward_url = get_model_url(args, "reward", "/v1/chat/completions") reward_output = await post(reward_url, { "model": "reward", - "messages": [{"role": "user", "content": sample.prompt + actor_output["text"]}], + "messages": [*prompt_messages, {"role": "assistant", "content": response_text}], }) # ... 处理输出并返回 Sample @@ -446,7 +471,7 @@ async def generate_with_models(args, sample, sampling_params): ### Q: 可以不训练，只用 `--vllm-config` 做推理吗？ -虽然 `--vllm-config` 是为 vime 的训练循环设计的，但你可以通过配置仅 rollout 的运行来实现纯推理场景。对于完全独立的 vLLM 推理服务，建议直接使用 vLLM 原生的 `launch_server`，或使用 `--rollout-external` 模式连接预部署的引擎。 +虽然 `--vllm-config` 是为 vime 的训练循环设计的，但你可以通过配置仅 rollout 的运行来实现纯推理场景。对于完全独立的 vLLM 推理服务，建议直接使用公开的 `vllm serve` 命令，或使用 `--rollout-external-engine-addrs` 连接预部署的引擎。 ### Q: `--vllm-config` 和 `--prefill-num-servers` 是什么关系？ diff --git a/docs/zh/developer_guide/ci.md b/docs/zh/developer_guide/ci.md index 8e35f606d..5fbd123fb 100644 --- a/docs/zh/developer_guide/ci.md +++ b/docs/zh/developer_guide/ci.md @@ -1,122 +1,45 @@ # CI（持续集成） -vime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发——给 PR 添加特定 label 即可运行对应的测试套件。 +Vime 使用 Buildkite 进行持续集成。提交到仓库中的 pipeline 是 +`.buildkite/pipeline.yml`。 -## 工作原理 +## 始终运行的检查 -工作流定义在 `.github/workflows/pr-test.yml`（由 `pr-test.yml.j2` 自动生成）。每个 CI 任务会： +每个 pull request 都会运行以下 CPU step： -1. 在自托管 GPU runner 上通过 `docker run` 运行；大多数测试使用 `inferactinc/public:vime-latest`，镜像验证使用 `inferactinc/public:vime-test-latest`。 -2. 通过 `pip install -e . --no-deps` 安装 vime。 -3. 通过 `tests/ci/gpu_lock_exec.py --count ` 获取所需数量的 GPU。 -4. 执行测试文件：`python .py` 或 `python tests/.py`。如果测试位于 `tests/plugin_contracts/` 这样的子目录，CI 也会自动处理。 +| Step | 覆盖范围 | +|---|---| +| `pre-commit` | 格式化、lint 与仓库规则 | +| `plugin-contracts` | customization contract 与 CPU 测试 | +| `agent-adapter` | agent adapter 行为 | +| `upstream-sync-cpu` | 从上游同步的 CPU 测试 | +| `utils` | `tests/utils` | -每个测试文件遵循统一的模式：`prepare()` 函数下载模型和数据集，`execute()` 函数构建命令行参数并调用 `U.execute_train(...)`。 +权威命令与队列配置位于 `.buildkite/pipeline.yml`。 -## CI Labels +## GPU 套件 -给 PR 添加 label 即可触发对应的测试套件： +CPU step 通过后，Buildkite build 会显示名为 `Run GPU test suites?` 的 +block step。可以选择一个或多个套件： -| Label | Job | 说明 | -|---|---|---| -| `run-ci-short` | `e2e-test-short` | Qwen2.5-0.5B 轻量级冒烟测试（4 GPU），用于快速反馈。 | -| `run-ci-megatron` | `e2e-test-megatron` | 核心 Megatron 训练测试，覆盖 Dense、MoE、PPO、MTP 等。 | -| `run-ci-precision` | `e2e-test-precision` | 数值精度校验（并行一致性检查）。 | -| `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint 保存/加载正确性（同步和异步保存）。 | -| `run-ci-image` | `e2e-test-image` | 在 `inferactinc/public:vime-test-latest` 镜像上运行**全部**测试（用于镜像验证）。 | -| `run-ci-changed` | `e2e-test-changed` | **动态**检测 PR 中新增或修改的测试文件，仅运行这些测试。 | +- `short` +- `vllm-config` +- `megatron` +- `vime-customized` +- `precision` +- `ckpt` -所有 label 也可通过 `workflow_dispatch`（在 Actions 页面手动触发）来运行。 +`.buildkite/gpu_suites.py` 会把所选套件展开为每个测试一个 Buildkite +job。验证 Dockerfile 或 vLLM patch 修改时，通过 `VIME_CI_IMAGE` 指定不可变的 +候选镜像 digest；未设置时使用 `vllm/vime:latest`，且 PR 合入前不得更新该标签。 -## 重点 Label 说明 +## 注册测试 -### `run-ci-changed` — 仅运行新增或修改的测试 +- 始终运行的 CPU 测试加入 `.buildkite/pipeline.yml` 中对应的命令。 +- GPU 测试加入 `.buildkite/gpu_suites.py` 中对应的套件，并同步更新 + `.buildkite/pipeline.yml` 显示的测试数量。 +- `.buildkite/README.md` 必须与 pipeline 行为保持一致。 -这是开发中最常用的 label。当你新增或修改了测试文件时，只需给 PR 添加 `run-ci-changed`，CI 会自动： - -1. **检测**相对于 `origin/main` 新增或修改的 `tests/test_*.py` 或 `tests/plugin_contracts/test_*.py` 文件（通过 `git diff --diff-filter=AM`）。 -2. **提取**每个测试文件中的 `NUM_GPUS` 值。 -3. **构建**动态 GitHub Actions matrix，并行运行每个测试。 - -这意味着你不需要手动在 workflow 中注册新测试——只需确保测试文件顶部有 `NUM_GPUS = ` 常量，`run-ci-changed` 就会自动识别并运行。 - -**示例**：如果你的 PR 新增了 `tests/test_mimo_7B_mtp_only_grad.py`（其中 `NUM_GPUS = 8`），添加 `run-ci-changed` label 后会自动在 8 张 GPU 上运行该测试。 - -### `run-ci-image` — 在测试镜像上运行全部测试 - -这会在 `inferactinc/public:vime-test-latest` Docker 镜像上运行**所有**已注册的测试。适用于： - -- 验证新构建的 Docker 镜像是否可用。 -- 在合并前做全面的测试检查。 - -由于包含所有测试，GPU 占用时间较长——日常开发请优先使用更有针对性的 label。 - -### `run-ci-megatron` — 核心 Megatron 测试 - -这是验证 Megatron 后端改动的主要 label，覆盖： - -- Dense 模型：GLM4-9B、Qwen3-4B（PPO） -- MoE 模型：Qwen3-30B-A3B（DeepEP）、Qwen3.6-35B-A3B PD + Mooncake、Moonlight-16B-A3B -- 特殊场景：MiMo-7B MTP、Qwen2.5-0.5B debug rollout-then-train - -所有测试使用 8 张 GPU。如果你正在修改 Megatron 训练逻辑、loss 计算或 checkpoint 转换，应该使用这个 label。 - -## 编写新测试 - -1. 创建 `tests/test_.py`，遵循标准模式： - -```python -import os -import vime.utils.external_utils.command_utils as U - -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" -NUM_GPUS = 4 # 此常量会被 run-ci-changed 自动读取 - -def prepare(): - U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - # 按需下载数据集 ... - -def execute(): - # 构建参数字符串并调用 U.execute_train(...) - ... - -if __name__ == "__main__": - prepare() - for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): - os.environ.pop(proxy_var, None) - execute() -``` - -2. **快速验证**：直接推送测试文件，给 PR 添加 `run-ci-changed` label，测试会被自动检测并运行。 - -3. **注册到固定 label 组**：编辑 `.github/workflows/pr-test.yml.j2`，在对应 job 的 `tests` 列表中添加条目，然后重新生成： - -```bash -cd .github/workflows && python generate_github_workflows.py -``` - -记得同时提交 `.j2` 和生成的 `.yml` 文件。 - -## Workflow 生成 - -工作流文件 `pr-test.yml` 是从 Jinja2 模板 `pr-test.yml.j2` 自动生成的。**不要直接编辑 `pr-test.yml`**。修改步骤： - -1. 编辑 `.github/workflows/pr-test.yml.j2`。 -2. 运行 `python .github/workflows/generate_github_workflows.py`。 -3. 同时提交两个文件。 - -## Customization 契约测试 - -如果你要运行通过函数路径加载的 customization hook 契约测试，可以使用： - -```bash -python -m pytest \ - tests/plugin_contracts/test_plugin_rollout_contracts.py \ - tests/plugin_contracts/test_plugin_generate_contracts.py \ - tests/plugin_contracts/test_plugin_path_loading_contracts.py \ - tests/plugin_contracts/test_plugin_runtime_hook_contracts.py -``` - -这些测试文件也支持直接执行 `python tests/plugin_contracts/.py`。它们声明了 `NUM_GPUS = 0`，因此可以被 `run-ci-changed` 自动识别，同时不会被当作 GPU 重型端到端测试。 +触发远程 Buildkite job 前，应先在本地运行完全相同的命令。GPU 测试失败 +时，先使用相同镜像和环境在 H200 节点复现并修复；本地通过后再重跑远程 +套件。 diff --git a/docs/zh/developer_guide/debug.md b/docs/zh/developer_guide/debug.md index c41a3a95c..96d4fbf84 100644 --- a/docs/zh/developer_guide/debug.md +++ b/docs/zh/developer_guide/debug.md @@ -48,6 +48,54 @@ vime 支持将训练部分和推理部分分开进行调试，从而实现： 开启后，会从 `args.load_debug_rollout_data.format(rollout_id=rollout_id)` 来加载数据，并且不会初始化 vllm（自动设置 `debug_train_only=True`）。可以以这种方式来固定训练部分的输入，对训练部分进行调优，例如切换各种并行。 +4. `--save-debug-train-data /your/saved/debug/train_{rollout_id}.pt` + + 每个 rollout 只保存一个训练侧文件。只有 Pipeline Parallel 最后一级和 Tensor Parallel rank 0 参与：跨 Context Parallel rank 逐个还原 `log_probs`、`ref_log_probs`、`values`、`advantages`、`returns`、`kl`、`entropy` 等 response-token tensor；Context Parallel rank 0 会将每个完整 tensor 立即搬到 CPU，避免它们在显存中累计，最后再把不同的 Data Parallel shard 汇总给一个 writer。 + + version 2 payload 对标 rollout debug dump：顶层 `samples` 列表每项是一个训练样本的 dict（含 `sample_index`、`data_parallel_rank` 以及 `tokens`、`log_probs`、`advantages` 等 per-sample 字段），并按 `sample_index` 排序，从而和 rollout dump 的 `samples` 一一对齐（用 `sample_index` ↔ rollout 侧的 `index` 来 join）。并列的 `dp_shards` key 保留 DP/micro-batch 排布——每项记录 `rank`、`data_parallel_rank`、该分片的 `sample_indices`，以及 DP-local 调度（`micro_batch_indices`、`num_microbatches`、`global_batch_sizes`）——且不重复存储任何 per-sample tensor。`raw_reward` 等整批字段在顶层只存一份。若某些样本没有 `sample_index`（自定义 rollout 新建 `Sample` 时会是 `None`），则 samples 保持 DP-gather 顺序并打印一条 warning。开启或关闭 CP 时，response-token 字段都是相同的完整 response 格式。在跳过 actor log-prob 单独重算的配置下（`can_reuse_log_probs_in_loss` 或 `--use-rollout-logprobs`），actor 的 `log_probs` 会直接从训练前向里快照下来（按 rollout position 归位，无额外前向），所以 dump 里依然会带上它。 + +## INT4 / Compressed-Tensors 量化 Checkpoint 问题 + +使用 INT4 量化模型（如 `compressed-tensors` 的 `W4A16`）时，checkpoint 的 `config.json` 中有一个 `quantization_config.ignore` 列表，指定哪些参数**不**做量化。在线权重更新（Megatron → vLLM）时，vime 也会读取这个 ignore list 来决定哪些参数需要 INT4 量化。ignore list 不正确会导致静默错误： + +1. **MoE 路由权重（`mlp.gate.weight`）变成全零** + + MoE 的路由权重（`mlp.gate.weight`，shape `[num_experts, hidden_size]`）是一个普通的 2D weight tensor，但它**不是** Linear 层的权重。如果它不在 ignore list 中，在线量化器会把它 INT4 量化为 `weight_packed`、`weight_scale`、`weight_zero_point` 等。然而 vLLM 不会以量化名称来加载路由权重，因此这些参数在 `load_weights` 时被静默跳过，导致 gate 权重全零。 + + **修复方法**：确保 `config.json` 的 ignore list 中包含 `"re:.*mlp\\.gate\\..*"`。 + +2. **其他非 Linear 的 2D 权重** + + 类似问题可能出现在任何不是真正 Linear 层的 2D `.weight` tensor 上，例如 `model.embed_tokens.weight`。务必检查 ignore list 覆盖了所有非 Linear 权重。 + + **推荐的 ignore 配置**（以 GLM 系 MoE 模型为例）： + ```json + "ignore": [ + "lm_head", + "model.embed_tokens.weight", + "re:.*self_attn.*", + "re:.*mlp\\.shared_experts.*", + "re:.*mlp\\.gate_up_proj.*", + "re:.*mlp\\.gate_proj.*", + "re:.*mlp\\.up_proj.*", + "re:.*mlp\\.down_proj.*", + "re:.*eh_proj.*", + "re:.*mlp\\.gate\\..*" + ] + ``` + +3. **safetensors 分片缺失** + + 转换工具偶尔可能产出不完整的 checkpoint（例如缺少 `model-00010-of-00093.safetensors`）。转换完成后，务必检查： + - `.safetensors` 文件数量是否与预期一致。 + - `model.safetensors.index.json` 中是否包含所有 layer 的条目。 + - 抽查关键 layer（如第一个 MoE layer）的 key 数量是否正确。 + +4. **如何排查** + + - 使用 `--check-weight-update-equal` 验证 Megatron → vLLM 权重同步后的值是否正确。如果某个参数在 vLLM 侧全为零，说明它可能被错误量化或在 checkpoint 中缺失。 + - 使用 `--debug-rollout-only` 配合少量 GPU，快速测试 vLLM 能否从量化 checkpoint 正常生成文本。 + ## Debug vllm illegal memory access (IMA) 在进行大规模 RL 时，不时会遇到 vLLM IMA 的问题，以下是我们的一些 debug 建议： diff --git a/docs/zh/developer_guide/install_flashqla.md b/docs/zh/developer_guide/install_flashqla.md index 647f74c53..1010af476 100644 --- a/docs/zh/developer_guide/install_flashqla.md +++ b/docs/zh/developer_guide/install_flashqla.md @@ -19,10 +19,11 @@ FlashQLA 是 Qwen GDN kernel 的可选运行后端。安装 FlashQLA 后，仍 ## Docker 镜像 -标准 CUDA Docker 镜像会默认安装 FlashQLA： +直接构建 Dockerfile 时默认不安装 FlashQLA；发布 target 会显式开启。手动构建时请传入相同参数： ```bash docker build \ -f docker/Dockerfile \ + --build-arg INSTALL_FLASHQLA=1 \ -t vime:flashqla . ``` diff --git a/docs/zh/developer_guide/profiling.md b/docs/zh/developer_guide/profiling.md index 7e43793f6..d96133f14 100644 --- a/docs/zh/developer_guide/profiling.md +++ b/docs/zh/developer_guide/profiling.md @@ -9,7 +9,7 @@ - 从日志确认router/worker地址 - start_profile - 发送少量推理请求 --（可选）stop_profile；或达到max_iterations后自动落盘 +-（可选）stop_profile；或达到max_iterations后自动写入 trace - 在torch_profiler_dir查看trace文件 @@ -44,21 +44,19 @@ vLLM只有在启动时配置了`--profiler-config`，才会注册`/start_profile |------|------| | `profiler` | `"torch"` 或 `"cuda"` | | `torch_profiler_dir` | trace输出目录（绝对路径） | -| `max_iterations` | worker记录超过N步后自动stop并落盘（条件为`> N`） | +| `max_iterations` | worker记录超过N步后自动stop并写入 trace（条件为`> N`） | | `ignore_frontend` | 建议`true`，仅profile worker，降低前端开销 | -**防止`stop_profile`时RPC超时：** vLLM APIServer与EngineCore/worker之间通过内部RPC通信。手动调用`stop_profile`触发trace落盘可能耗时数分钟，而默认`VLLM_RPC_TIMEOUT`仅**10秒**（10000 ms），容易导致flush中断或trace不完整。Profiling时建议设为**30分钟**（1800000 ms）。 +`/stop_profile` 会等待 EngineCore 和全部 worker 完成 trace 写盘，因此较大的 +profile 可能需要数分钟；trace 文件写完前不要中断该请求。 -该变量须在**启动train、拉起vLLM之前**传入Ray worker环境（仅在本机shell `export`不一定会进入Ray job）。在`ray job submit`的`runtime-env-json`中写入，例如： +常规 worker 环境仍通过 `ray job submit` 的 `runtime-env-json` 传入，例如： ```bash -export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" } }" @@ -111,7 +109,7 @@ python tools/profile_rollout.py \ ### 停止Profiling（可选） -若在`--vllm-profiler-config`中设置了`max_iterations`，worker在记录足够步数后会**自动stop并落盘**，实践中发完推理后常可直接在`torch_profiler_dir`看到trace，**不必**再手动`stop_profile`。需要提前结束采集时再执行： +若在`--vllm-profiler-config`中设置了`max_iterations`，worker在记录足够步数后会**自动stop并写入 trace**，实践中发完推理后常可直接在`torch_profiler_dir`看到trace，**不必**再手动`stop_profile`。需要提前结束采集时再执行： ```bash python tools/profile_rollout.py \ @@ -124,8 +122,8 @@ python tools/profile_rollout.py \ 在sleep_rollout等待期间，执行步骤如下： 1. `profile_rollout.py --action start` -2. 向router或**直连worker**发送少量completion请求（2～4条即可，trace会很大） -3. （可选）`profile_rollout.py --action stop`；或等待`max_iterations`触发自动落盘 +2. 向router或**直连worker**发送少量completion请求（通常2～4条即可，trace会很大） +3. 如果依赖自动写入 trace，要注意 `max_iterations` 的停止条件是 `> N`。例如 `max_iterations=3` 时，需要发 4 条请求；否则请手动执行 `profile_rollout.py --action stop` 4. 在`torch_profiler_dir`查看trace 请求示例（`model`使用HF checkpoint路径）： @@ -162,9 +160,9 @@ python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-rank | 现象 | 处理 | |------|------| | `POST /start_profile` 404 | 用JSON传`--vllm-profiler-config`；重启job | -| start成功但目录为空 | 确认curl打到worker且返回200；适当增大`max_iterations`或补发推理 | +| start成功但目录为空 | 确认curl打到worker且返回200；若 `max_iterations=3`，请发 4 条请求，或手动执行 `stop_profile` | | router 503 | 确认当前job的router端口；改直连worker | -| stop很慢或超时 | 增大`VLLM_RPC_TIMEOUT`；减少请求条数 | +| stop 很慢 | 等待 trace 写盘完成；减少请求条数 | ## 8. 完整可运行示例 @@ -212,13 +210,10 @@ launch_train_for_profiling() { source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" - export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" } }" @@ -290,16 +285,15 @@ run_profiling_session() { echo "=== 1/3 start_profile (all workers via router) ===" python tools/profile_rollout.py --router-url "${router_url}" --action start - echo "=== 2/3 send completions (direct to worker; 3 requests) ===" - for i in 1 2 3; do - curl -sS -X POST "${worker_url}/v1/completions" \ + echo "=== 2/3 send completions (direct to worker; 4 requests so max_iterations=3 can auto-flush) ===" + for i in 1 2 3 4; do + response="$(curl -sS -X POST "${worker_url}/v1/completions" \ -H "Content-Type: application/json" \ - -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ - | head -c 400 - echo + -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}")" + printf '%s\n' "${response:0:400}" done - echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" + echo "=== 3/3 list trace files (max_iterations=3 auto-stop uses > N; add --action stop if needed) ===" sleep 2 find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" diff --git a/docs/zh/developer_guide/trace.md b/docs/zh/developer_guide/trace.md index 99a16d042..b7abf2bdc 100644 --- a/docs/zh/developer_guide/trace.md +++ b/docs/zh/developer_guide/trace.md @@ -41,7 +41,7 @@ python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt ## 给自定义代码打点 -在自定义 rollout 或 reward 逻辑中——包括 agentic workflow 里的 agent step、tool call、sandbox 执行、verifier 调用等——可以直接复用 `vime.utils.trace_utils` 里的工具： +在自定义 rollout 或 reward 逻辑中——包括 agentic workflow 里的 agent step、tool call、sandbox 执行、verifier 调用等——可以直接复用 `vime.observability.trace_utils` 里的工具： - `trace_span(target, name, attrs=...)`：记录一段持续时间。 - `trace_event(target, name, attrs=...)`：记录一个瞬时事件。 @@ -57,7 +57,7 @@ python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt vime 主 rollout 流程里就是这样用的。例如 `generate_and_rm(...)` 按 sample 打点，而 `generate_and_rm_group(...)` 按 group 打点： ```python -from vime.utils.trace_utils import trace_function +from vime.observability.trace_utils import trace_function @trace_function("generate_and_rm", target="sample") @@ -104,7 +104,7 @@ async def custom_rollout_batch(samples, **kwargs): 如果想统一记录 vLLM 返回的 generation 元信息，可以复用 `build_vllm_meta_trace_attrs`： ```python -from vime.utils.trace_utils import build_vllm_meta_trace_attrs, trace_span +from vime.observability.trace_utils import build_vllm_meta_trace_attrs, trace_span with trace_span(sample, "vllm_generate") as span: output = await post(url, payload) @@ -116,4 +116,3 @@ with trace_span(sample, "vllm_generate") as span: - 先保存少量 rollout；单个 dump 的 sample 数量适中时，viewer 会更容易阅读。 - viewer 直接基于保存下来的 `.pt` dump 工作，因此可以把文件拷到别的机器离线分析。 - 如果你想看的是 vLLM 自身的 GPU / kernel 级 profiling trace，请参考 [性能分析](./profiling.md)。 - diff --git a/docs/zh/examples/deepseek-r1.md b/docs/zh/examples/deepseek-r1.md new file mode 100644 index 000000000..4c30dd904 --- /dev/null +++ b/docs/zh/examples/deepseek-r1.md @@ -0,0 +1,206 @@ +# 128xH100 训练 DeepSeek R1 + +这里是使用 128xH100 进行 DeepSeek R1 RL 训练的示例。 + +我们会使用 bf16 进行训练，128x128 blockwise quant 的 fp8 格式进行推理，模型最大回复长度为 32k，并训练中会使用 dynamic sampling 对数据进行筛选。 + +在并行上，vLLM 方面我们会开启专家并行（`--vllm-enable-expert-parallel`）与数据并行（`--vllm-data-parallel-size 8`），DeepEP 默认关闭；megatron 部分我们采用 tp8、pp4、ep32、cp4。 + +⚠️ 为了节省 GPU 显存，我们会使用 CPU Adam，每个 node（8xH100）会占用 1.4~1.5B 内存。如果单机的内存不够，可以通过增加 GPU，扩大并行的方式解决。 + +## 环境准备 + +搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。 + +准备 DeepSeek R1 的 ckpt 首先需要在多机均可访问到的地址（下记为 `$BASE_DIR`）上下载 DeepSeek-R1： + +```bash +hf download deepseek-ai/DeepSeek-R1 --local-dir $BASE_DIR/DeepSeek-R1 +``` + +DeepSeek-R1 的 huggingface ckpt 为 block-quant 的 fp8 格式，为了转换一个 Megatron 可以加载的 torch dist 格式，需要先转化一个 bf16 的 huggingface ckpt： + +```bash +cd vime/ +python tools/fp8_cast_bf16.py --input-fp8-hf-path $BASE_DIR/DeepSeek-R1 --output-bf16-hf-path $BASE_DIR/DeepSeek-R1-bf16/ +``` + +之后我们需要将 bf16 版本的 DeepSeek-R1 转换为 torch dist 格式。具体为在 4 台机器上分别执行： + +```bash +cd vime/ +source scripts/models/deepseek-v3.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=4 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 8 \ + --expert-tensor-parallel-size 1 \ + --expert-model-parallel-size 4 \ + --decoder-first-pipeline-num-layers 7 \ + --decoder-last-pipeline-num-layers 6 \ + --hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ \ + --save $BASE_DIR/DeepSeek-R1_torch_dist/ +``` + +其中 `MASTER_ADDR` 为 node0 的 ip，`NODE_RANK` 表示这是第几台机器，这两者就像是在多机 `torchrun` 的时候进行的配置。 + +## 执行训练 + +在 node0 运行： + +```bash +cd vime/ +bash scripts/run-deepseek-r1.sh +``` + +在其他 node 需要通过如下的指令加入 ray 集群： + +```bash +ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" +``` + +或者如果你能获取到所有节点的 ip 列表，例如有一个 mpi hostfie（每一行为 `ip slot=8`），那么可以在 `scripts/run-deepseek-r1.sh` 中的 `ray start --head` 指令之后加入如下的指令，从而只需要从 node0 执行训练： + +```bash +for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do + if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 vllm ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" & +done +wait +``` + +### 参数简介 + +```bash +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/deepseek-v3.sh" +``` + +从 [scripts/models/deepseek-v3.sh](https://github.com/vllm-project/vime/blob/main/scripts/models/deepseek-v3.sh) 读取模型的 config。这些 config 都是 megatron 的参数。在使用 megatron 进行训练的时候，megatron 无法从 ckpt 中读取模型 config，需要我们自行配置。我们在 [scripts/models](https://github.com/vllm-project/vime/tree/main/scripts/models/) 中提供了一些样例。 + +#### CKPT_ARGS + +```bash +CKPT_ARGS=( + # vllm 需要的 hf ckpt，我们也会从这里读 tokenizer + --hf-checkpoint $BASE_DIR/DeepSeek-R1/ + #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ + --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/ + # actor 的 load dir，如果是空的，会从 `ref_load` 里面读 + --load $BASE_DIR/DeepSeek-R1_vime/ + --save $BASE_DIR/DeepSeek-R1_vime/ + --save-interval 20 +) +``` + +vime 会根据 `hf_checkpoint` 中的量化配置从而在训练中进行在线量化。例如当前的例子中，我们使用的是 DeepSeek R1 的 fp8 ckpt，那么在进行参数更新的时候，我们会首先将参数进行 blockwise quant，再传至 vllm。 + +#### PERF_ARGS + +一堆 megatron 的并行参数，只有 `--use-dynamic-batch-size` 与 `--max-tokens-per-gpu` 是 vime 添加的。 + +megatron 的部分，我们配置了 tp8、pp4、cp4、ep32，由于 DeepSeek-R1 有 61 层，不能被 4 整除，所以我们专门配置最后一个 pp stage 为 13 层。 + +`max_tokens_per_gpu` 是指每张卡最多跑多少 token，在开启 `use_dynamic_batch_size` 之后，会尽可能将一个 batch 内部长短不一的数据拼到 `max_tokens_per_gpu`，从而组成动态的 micro batch size，如果有一条数据长度超过了 `max_tokens_per_gpu`，则自成一条，不会对数据进行截断。在开启 context parallel (CP) 时，会让 CP 张卡去上的数据去共享总长为 `CP * max_tokens_per_gpu` 的 token。 + +在开启 dynamic_batch_size，会忽略传统的 `micro_batch_size`。 + +⚠️ vime 总是会通过 data packing 的方法训练模型，并且严格保证 per sample loss 或 per token loss，也就是开启 dynamic batch size 不会对 loss 计算有影响，推荐开启。 + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 13 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) +``` + +#### GRPO_ARGS + +目前 vime 这是一些 grpo 相关的参数： + +```bash +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) +``` + +如果希望训练时不加载 reference model，需要去掉 `--use-kl-loss` 并设置 `--kl-coef 0.00`（默认值为 0）。 + +#### OPTIMIZER_ARGS + +我们通过了如下几个参数配置了 CPU Adam，用来节省显存。 + +```bash +OPTIMIZER_ARGS=( + ... + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) +``` + +#### VLLM_ARGS + +这些是 vLLM 所需的参数。`--rollout-num-gpus-per-engine` 表示单个 engine 的 worker GPU 总数；这里它等于 `tensor_parallel_size * data_parallel_size`，而不只是 tensor-parallel size。其他 vLLM 参数通过添加 `--vllm-` 前缀传给 vime。为了充分利用 vLLM 的大 EP 推理能力，我们通过 `--vllm-enable-expert-parallel` 开启专家并行，通过 `--vllm-data-parallel-size 8` 开启 DP attention。DeepEP 默认关闭，可通过脚本中注释掉的 flag 开启。 + +最后的 `--vllm-server-concurrency` 是 vime 的特有参数，是为了防止同时发给 vllm server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.7 + --vllm-enable-expert-parallel + + # dp attention + --vllm-data-parallel-size 8 + + # enable deepep for vllm + + # mtp + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +#### MISC_ARGS + +一些额外的 megatron 配置。注意这里配置了 megatron 的 deepep。 + +```bash +MISC_ARGS=( + ... + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) +``` diff --git a/docs/zh/examples/glm4-9B.md b/docs/zh/examples/glm4-9B.md new file mode 100644 index 000000000..bca98aab2 --- /dev/null +++ b/docs/zh/examples/glm4-9B.md @@ -0,0 +1,278 @@ +# 8xH100 训练 GLM4-9B + +## 环境准备 + +拉取 `vllm/vime:latest` 镜像后，用如下方式初始化镜像环境： + +```bash +cd /root/ +git clone https://github.com/vllm-project/vime.git +cd vime/ +pip install -e . --no-deps +``` + +下载模型与数据： + +```bash +# hf checkpoint +hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 + +# train data +hf download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +# eval data +hf download --repo-type dataset zhuzilin/aime-2024 \ + --local-dir /root/aime-2024 +``` + +将 huggingface checkpoint 转换成 megatron 可以加载的 huggingface checkpoint： + +```bash +# mcore checkpoint +cd /root/vime +source scripts/models/glm4-9B.sh +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/GLM-Z1-9B-0414 \ + --save /root/GLM-Z1-9B-0414_torch_dist +``` + +## 执行训练 + +执行训练： + +```bash +cd /root/vime +bash scripts/run-glm4-9B.sh +``` + +### 参数简介 + +这里我们简单介绍一下脚本 [run-glm4-9B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-glm4-9B.sh) 中的各个组成部分： + +#### MODEL_ARGS + +```bash +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/glm4-9B.sh" +``` + +从 [scripts/models/glm4-9B.sh](https://github.com/vllm-project/vime/blob/main/scripts/models/glm4-9B.sh) 读取模型的 config。这些 config 都是 megatron 的参数。在使用 megatron 进行训练的时候，megatron 无法从 ckpt 中读取模型 config，需要我们自行配置。我们在 [scripts/models](https://github.com/vllm-project/vime/tree/main/scripts/models/) 中提供了一些样例。 + +⚠️ 注意检查模型文件中的 `--rotary-base` 等配置是否对应你当前训练模型的配置，因为同一个模型结构的不同模型可能有不同的取值。在这种情况下，你可以在导入模型参数后在脚本里进行覆盖，例如： + +```bash +source "${SCRIPT_DIR}/models/glm4-9B.sh" + +MODEL_ARGS += ( --rotary-base 10000 ) +``` + +#### CKPT_ARGS + +```bash +CKPT_ARGS=( + # vllm 需要的 hf ckpt，我们也会从这里读 tokenizer + --hf-checkpoint /root/GLM-Z1-9B-0414 + # reference model 的 ckp + --ref-load /root/GLM-Z1-9B-0414_torch_dist + # actor 的 load dir，如果是空的，会从 `ref_load` 里面读 + --load /root/GLM-Z1-9B-0414_vime/ + --save /root/GLM-Z1-9B-0414_vime/ + --save-interval 20 +) +``` + +#### ROLLOUT_ARGS + +```bash +ROLLOUT_ARGS=( + # prompt 数据集，每行是个 json + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + # 如果 prompt 的 `input_key` 中是 openai message， + # 会进行 tokenizer.apply_chat_template(...) + --apply-chat-template + # 是否 shuffle 数据 + --rollout-shuffle + + # reward model 类型， + # vime 提供了很多类型以及用于自定义的 --custom-rm-path + --rm-type deepscaler + + # 一共要训练多少 rollout + --num-rollout 3000 + # 一个 rollout 有多少 prompt + --rollout-batch-size 32 + # 每个 prompt 采多少回复 + # 一个 rollout 会有 rollout_batch_size * n_samples_per_prompt 条 + --n-samples-per-prompt 8 + # rollout sampling param + --rollout-max-response-len 8192 + --rollout-temperature 0.8 + + # 一次 rollout 对应几个训练步 + --num-steps-per-rollout 1 + # 是否在训练时 balance data，可能对速度有好处 + --balance-data +) +``` + +#### EVAL_ARGS + +eval 的时候基本上是会继承所有 rollout 的参数，但是我们提供了一些可以 rollout 配置覆盖的参数，从而实现训练和 eval 用不同的采样策略。 + +```bash +EVAL_ARGS=( + --eval-interval 5 + --eval-prompt-data /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 0.7 +) +``` + +#### PERF_ARGS + +一堆 megatron 的并行参数，只有 `--use-dynamic-batch-size` 与 `--max-tokens-per-gpu` 是 vime 添加的。 + +`max_tokens_per_gpu` 是指每张卡最多跑多少 token，在开启 `use_dynamic_batch_size` 之后，会尽可能将一个 batch 内部长短不一的数据拼到 `max_tokens_per_gpu`，从而组成动态的 micro batch size，如果有一条数据长度超过了 `max_tokens_per_gpu`，则自成一条，不会对数据进行截断。在开启 context parallel (CP) 时，会让 CP 张卡去上的数据去共享总长为 `CP * max_tokens_per_gpu` 的 token。 + +在开启 dynamic_batch_size，会忽略传统的 `micro_batch_size`。 + +⚠️ vime 总是会通过 data packing 的方法训练模型，并且严格保证 per sample loss 或 per token loss，也就是开启 dynamic batch size 不会对 loss 计算有影响，推荐开启。 + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 2 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 4608 +) +``` + +#### GRPO_ARGS + +目前 vime 这是一些 grpo 相关的参数： + +```bash +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) +``` + +#### OPTIMIZER_ARGS + +```bash +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) +``` + +#### VLLM_ARGS + +这些是 vLLM 所需的参数。在默认并行配置下，`--rollout-num-gpus-per-engine` 对应 vLLM 的 `tensor_parallel_size`；其他 vLLM 参数通过添加 `--vllm-` 前缀传给 vime。 + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 +) +``` + +### 训推一体 + +在原始的脚本中，资源配置如下： + +```bash +ray job submit ... \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + ... +``` + +即开启训推分离，并且训练部分会使用 1 机 8 卡，推理会和训练共同使用这 8 张卡张卡。 + +如果想使用训推一体（colocate）的功能，需要加上 `--colocate` 并去掉 `--rollout-num-gpus`： + +```bash +ray job submit ... \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ... +``` + +此时，训练和推理就会共用这 8 张卡了。 + +⚠️ 在训推一体的训练时，megatron 始终会占据一些显存，所以需要通过调整 `--vllm-gpu-memory-utilization` 来降低 vllm 占据的显存比例。 + +### dynamic sampling + +vime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.github.io/) 中的 dynamic sampling。如果要开启 dynamic sampling，需要配置： + +```bash + --over-sampling-batch-size ${OVER_SAMPLING_BS} \ + --dynamic-sampling-filter-path \ + vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ +``` + +这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为： + +```bash + --rollout-batch-size 32 \ + --n-samples-per-prompt 8 \ + --over-sampling-batch-size 64 \ +``` + +那么 sampling 会直接采样 64 条 prompt，每条 prompt 采样 8 次。因为 vime 内部进行的是异步采样，所以我们会先后获得每个 prompt 的 8 条回复。在收到回复时，会用 `dynamic_sampling_filter_path` 对应的函数进行筛选，如果通过，则留下这 8 条数据，否则则丢掉。例子中的函数是判断回答是否全对或全错： + +```python +def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): + rewards = [sample.reward for sample in samples] + return torch.tensor(rewards, dtype=torch.float).std() > 0.0 +``` + +当我们收到了 32 * 8 条数据时，我们会立刻停止采样，而不会等剩余的数据采样完成。如果删除的数据超过了 32 条 prompt（剩余的小于 32 条 prompt），那么我们会再采样 64 条 prompt。 + +### partial rollout + +在进行 dynamic sampling 的过程中，会提前终止（abort）大量请求，我们可以通过配置 `--partial-rollout` 参数来将生成到一半的请求保存至 data buffer，在下一个 rollout 中取出来继续进行数据生成，从而进一步优化性能。 + +可以通过配置 `--buffer-filter-path` 来自定义如何从 buffer 中取出数据，默认的函数为： + +```python +def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]: + num_to_pop = min(len(buffer), num_samples) + samples = buffer[:num_to_pop] + del buffer[:num_to_pop] + return samples +``` + +即每次取出前 `num_samples` 个 prompt 对应的 `num_samples * n_samples_per_prompt` 条数据。 + +⚠️ 每条 partial rollout sample 的 `sample.metadata` 中存储了第一次进行生成的 rollout id，可以用于数据过滤。 diff --git a/docs/zh/examples/glm4.7-30B-A3B.md b/docs/zh/examples/glm4.7-30B-A3B.md new file mode 100644 index 000000000..a787fcfd5 --- /dev/null +++ b/docs/zh/examples/glm4.7-30B-A3B.md @@ -0,0 +1,141 @@ +# 8×H100 训练 GLM-4.7-Flash + +## 环境准备 + +搭建环境、数据与 ckpt 转换均与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分转换为 GLM-4.7-Flash 即可。 + +### 下载模型 + +```bash +hf download THUDM/GLM-4.7-Flash --local-dir /root/GLM-4.7-Flash +``` +### 转换 Checkpoint + +可以用如下方法把 Hugging Face checkpoint 转化为 torch_dist 格式： + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm4.7-30B-A3B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/GLM-4.7-Flash/ \ + --save /root/GLM-4.7-Flash_torch_dist/ +``` + +## 执行训练 + +执行训练： + +```bash +cd /root/vime +bash scripts/run-glm4.7-30B-A3B.sh +``` + +### 参数简介 + +这里我们简单介绍一下脚本 [run-glm4.7-30B-A3B.sh](../../../scripts/run-glm4.7-30B-A3B.sh) 中的关键部分。 + +#### MoE 配置 + +GLM-4.7-Flash 是一个 MoE（混合专家）模型，包含 64 个路由专家（top-4 激活）和 1 个共享专家。共 47 层：1 层 dense 层 + 46 层 MoE 层。 + +1. 为了支持在 8×H100 环境中运行 GLM-4.7-Flash，我们需要开启 Megatron 的 CPU Adam 以节省显存： + + ```bash + OPTIMIZER_ARGS=( + ... + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + ) + ``` + +2. 开启 Megatron 支持的 MoE 优化，单机 8×H100 配置为 TP=1, EP=8： + + ```bash + PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + ... + ) + ``` + +3. 开启 vLLM 数据并行（`--vllm-data-parallel-size`）以提升 MoE 推理吞吐： + + ```bash + VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 8 + --vllm-enable-expert-parallel + ... + ) + ``` + +#### MTP 投机解码（推理加速） + +GLM-4.7-Flash 包含 1 层 MTP（Multi-Token Prediction）层，可用于推理时的投机解码来加速 rollout 生成。要启用此功能，在 `VLLM_ARGS` 中添加以下配置： + +```bash +VLLM_ARGS=( + ... + # MTP 投机解码 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +这会让 vLLM 使用模型的 MTP 层进行投机解码。MTP 层预测多个未来 token，vLLM 并行验证它们，从而加速生成。 + +> ⚠️ **注意**：投机解码会占用额外的 GPU 显存。如果遇到 OOM 问题，可以尝试降低 `--vllm-gpu-memory-utilization` 或关闭投机解码。 + +#### MTP 训练 + +vime 也支持将 MTP 层与主模型联合训练，适用于已实现 MTP 权重转换的模型（如 MiMo、GLM-4.7）。启用时，相关参数如下： + +```bash +# 在模型配置中添加 MTP 层数 +MODEL_ARGS+=(--mtp-num-layers 1) + +# 启用 MTP 训练 +SPEC_ARGS=( + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) +``` + +- `--mtp-num-layers 1`：告知 Megatron 从 checkpoint 中加载 MTP 层。 +- `--enable-mtp-training`：启用 MTP 层的梯度计算。不设置此标志时，MTP 层会被加载但冻结。 +- `--mtp-loss-scaling-factor 0.2`：MTP loss 相对于主策略 loss 的权重，默认为 0.2。 + +> **注意**：原生 DeepSeek 布局 loader 会按模型配置的层数映射 MTP 权重，包括 47 层的 GLM-4.7-Flash。 +> +> 对于其他支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。 + +### 多机适配 + +仓库中的 `scripts/run-glm4.7-30B-A3B.sh` 会启动本地单节点 Ray，并固定传入 `--actor-num-nodes 1`，不能直接作为多机启动器。要把这份配置适配到多机训练（例如 2×8 H100），需要先让所有 worker 加入同一个 Ray 集群，并修改启动脚本： + +对于多机环境，需要进行如下修改： + +- 将训练模型、数据放在所有机器都可以访问到的路径上； +- 设置各台机器都可以访问到的 `MASTER_ADDR`； +- 把 `--actor-num-nodes` 从 `1` 改为训练节点数； +- 去掉 CPU Adam 相关的配置，因为使用了 distributed optimizer，多机环境下 optimizer 的显存占比会明显下降。 +- 调整并行度：例如 TP=4, PP=2, EP=8, CP=2。 + +当总卡数并不能被 expert 总数（64）乘除时，可以使用 `--vllm-eplb-config` 来增加冗余的 expert。例如对于 24 卡的场景： + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 24 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 3 + --vllm-enable-expert-parallel + --vllm-eplb-config '{"num_redundant_experts": 16}' +) +``` diff --git a/docs/zh/examples/glm4.7-355B-A32B.md b/docs/zh/examples/glm4.7-355B-A32B.md new file mode 100644 index 000000000..52e519977 --- /dev/null +++ b/docs/zh/examples/glm4.7-355B-A32B.md @@ -0,0 +1,174 @@ +# 64xH100 训练 GLM-4.7 + +## 环境准备 + +搭建环境与下载数据的方法与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分替换为 GLM-4.7 即可。 + +### 前置条件 + +GLM-4.7 使用 vime 标准 Docker 环境即可。多机启动前，请确保所有机器都能访问同一个 `$BASE_DIR` 路径，并在启动 Ray worker 前先取消代理： + +```bash +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY +``` + +### 下载模型 + +```bash +hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B +``` + +### 转换 Checkpoint + +可以用如下方法把 Hugging Face checkpoint 转换为 torch_dist 格式（2 机 x 8 卡）： + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm4.5-355B-A32B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=2 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \ + --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/ +``` + +其中 `MASTER_ADDR` 是 node0 的 IP，`NODE_RANK` 表示当前机器的编号，配置方式与普通多机 `torchrun` 一致。 + +## 执行训练 + +从 node0 执行训练脚本： + +```bash +cd /root/vime +export BASE_DIR=/shared/path # 所有节点都能访问的共享路径 +bash scripts/run-glm4.7-355B-A32B.sh +``` + +### 参数简介 + +这里我们简单介绍一下 [run-glm4.7-355B-A32B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-glm4.7-355B-A32B.sh) 中的关键部分。 + +#### MoE 配置 + +GLM-4.7 是一个 MoE（混合专家）模型，包含 160 个路由专家（top-8 激活）和共享专家。模型共 92 层：3 层 dense + 89 层 MoE。 + +1. 为了支持在 64xH100 环境中运行 GLM-4.7，我们开启 Megatron 的 CPU Adam 来节省显存： + + ```bash + OPTIMIZER_ARGS=( + ... + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + ) + ``` + +2. 在 Megatron 中开启 MoE 优化。当前 64xH100 示例使用 TP=8、PP=4、CP=2、EP=16： + + ```bash + PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 2 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + ... + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 + ) + ``` + +3. 在 vLLM 中开启带 DP attention 的 MoE 优化： + + ```bash + VLLM_ARGS=( + --rollout-num-gpus-per-engine 32 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 4 + --vllm-enable-expert-parallel + ... + ) + ``` + +#### MTP 投机解码（推理加速） + +GLM-4.7 包含 MTP（Multi-Token Prediction）层，可以在推理阶段用于投机解码，加速 rollout 生成。启用方法是在 `VLLM_ARGS` 中加入： + +```bash +VLLM_ARGS=( + ... + # MTP 投机解码 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) +``` + +这样 vLLM 就会使用模型自带的 MTP 层进行投机解码。 + +> ⚠️ **注意**：投机解码会额外占用 GPU 显存。如果遇到 OOM，可以尝试降低 `--vllm-gpu-memory-utilization` 或暂时关闭投机解码。 + +#### MTP 训练 + +vime 也支持在 GLM-4.7 上将 MTP 层与主模型联合训练。启用时，相关参数如下： + +```bash +# 在模型配置中添加 MTP 层数 +MODEL_ARGS+=(--mtp-num-layers 1) + +# 启用 MTP 训练 +MTP_ARGS=( + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) +``` + +- `--mtp-num-layers 1`：告知 Megatron 从 checkpoint 中加载 MTP 层。 +- `--enable-mtp-training`：启用 MTP 层的梯度计算；不设置时 MTP 层会被加载但保持冻结。 +- `--mtp-loss-scaling-factor 0.2`：MTP loss 相对主策略 loss 的权重，默认值为 0.2。 + +> **注意**：`vime/backends/megatron_utils/hf_to_megatron/glm.py` 中的原生 loader 会同时映射普通层和 MTP 层权重。 + +#### 多机支持 + +这个示例本身就是多机训练配置。启动前请确认： + +- 模型权重和数据集放在所有节点都能访问到的路径； +- `MASTER_ADDR` 设置为所有节点都能访问到的地址； +- 在启动 Ray worker 前先取消代理； +- 提供一个 `HOSTFILE` 列出 worker IP（每行一个），并在启动前 `export HOSTFILE=/path/to/hostfile`； +- 并行度需要成套调整。默认示例使用 TP=8、PP=4、EP=16、CP=2，rollout 侧则使用 32 张卡 / engine + vLLM DP attention。 + +如果 rollout GPU 数与 expert 数（160）之间不能整除，可以通过 `--vllm-eplb-config` 增加冗余 expert。 + +## FP8 Rollout + +开源版 GLM-4.7 的 FP8 checkpoint 使用的是 per-channel 量化，目前无法在 vLLM 中直接启用 DeepEP。可以利用 vime 自带工具将其转换为 128x128 的 per-block FP8 checkpoint： + +```bash +cd /root/vime +python tools/convert_hf_to_fp8.py \ + --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \ + --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \ + --strategy block --block-size 128 128 \ + --max-workers 4 +``` + +随后把 `--hf-checkpoint` 改成 `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` 即可开启 FP8 rollout。 + +一个可参考的 FP8 `VLLM_ARGS` 配置如下： + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 32 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 32 + --vllm-enable-expert-parallel + --vllm-cudagraph-capture-sizes 5 10 20 40 $(seq 80 40 640) + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' + --vllm-all2all-backend deepep_high_throughput +) +``` diff --git a/docs/zh/examples/glm5.2-744B-A40B.md b/docs/zh/examples/glm5.2-744B-A40B.md new file mode 100644 index 000000000..f07daad3a --- /dev/null +++ b/docs/zh/examples/glm5.2-744B-A40B.md @@ -0,0 +1,166 @@ +# 256xH100 训练 GLM-5.2 744B-A40B + +这里是使用 32 节点、256 张 H100 训练 [GLM-5.2](https://z.ai/blog/glm-5.2) 的推荐配置示例。 + +这个配置使用 GLM-5.2 的 BF16 checkpoint 做 Megatron 训练，使用 FP8 checkpoint 做 vLLM rollout。下面假设 Hugging Face 上会提供两个地址： + +- BF16: `zai-org/GLM-5.2` +- FP8: `zai-org/GLM-5.2-FP8` + +## 环境准备 + +搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。多机启动前，请确保所有节点都能访问同一个 `$BASE_DIR` 路径。 + +### 下载模型 + +```bash +hf download zai-org/GLM-5.2 --local-dir $BASE_DIR/GLM-5.2 +hf download zai-org/GLM-5.2-FP8 --local-dir $BASE_DIR/GLM-5.2-FP8 +``` + +开源 GLM-5.2 的 config 使用 `model_type: glm_moe_dsa`，vime 将其映射到原生 DeepSeek-V3.2 loader，因为两者共享相同的 DSA 权重布局。 + +### 转换 Checkpoint + +训练侧需要把 BF16 Hugging Face checkpoint 转换为 Megatron 可加载的 torch_dist 格式。torch_dist 格式支持重新切分，所以转换时的并行布局**不需要**与训练一致；我们使用一个能满足 Megatron expert group 约束（在转换的节点数下成立）的布局即可。 + +可以在 4 台机器 / 32 卡上分别执行： + +```bash +cd /root/vime +pip install -e . --no-deps +source scripts/models/glm5.2-744B-A40B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun \ + --nproc-per-node 8 \ + --master-addr ${MASTER_ADDR} --master-port 12345 \ + --nnodes=4 --node-rank ${NODE_RANK} \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --tensor-model-parallel-size 8 \ + --pipeline-model-parallel-size 2 \ + --decoder-last-pipeline-num-layers 40 \ + --expert-model-parallel-size 16 \ + --expert-tensor-parallel-size 1 \ + --hf-checkpoint $BASE_DIR/GLM-5.2/ \ + --save $BASE_DIR/GLM-5.2_torch_dist/ +``` + +其中 `MASTER_ADDR` 是 node0 的 IP，`NODE_RANK` 表示当前机器编号。 + +`MODEL_ARGS` 里包含 `--allgather-cp` 这个 vime 自定义参数，所以 `tools/convert_hf_to_torch_dist.py` 也注册了它（转换时是 no-op）。在 32 卡上，Megatron 要求 `expert_tp(1) * expert_model_parallel * pp` 能整除 world size，因此转换时使用 `EP=16`（`1*16*2=32`）。由于 torch_dist 支持重新切分，转换出的 checkpoint 仍然可以在训练时以 `EP=32` 加载。 + +## 执行训练 + +从 node0 执行： + +```bash +cd /root/vime +export BASE_DIR=/shared/path +export MASTER_ADDR= +export HOSTFILE=$BASE_DIR/hostfile # 每行一个 worker IP，共 32 个节点 +bash scripts/run-glm5.2-744B-A40B.sh +``` + +如果不设置 `HOSTFILE`，需要手动在其他节点加入 Ray 集群。 + +### 参数简介 + +#### 模型配置 + +`scripts/models/glm5.2-744B-A40B.sh` 使用 GLM-5.2 的 DSA + cross-layer index sharing 配置：256 个 routed experts、top-8 激活、1 个 shared expert，模型共 78 层（3 层 dense + 75 层 MoE）。 + +DSA index sharing 的 schedule（例如 `index_topk_freq=4`、`index_skip_topk_offset=3`）从 Hugging Face config 中读取。Megatron 侧使用共享的 `vime_plugins.models.glm5.glm5:get_glm5_spec` provider，并开启： + +```bash +--allgather-cp +``` + +这会让 DSA + context parallel 使用 allgather-CP layout，并在 index-share provider 中对 index K/V 做 CP group gather。 + +#### 训练并行 + +默认脚本按 32 节点 256 卡配置： + +```bash +PERF_ARGS=( + --tensor-model-parallel-size 4 + --pipeline-model-parallel-size 8 + --decoder-first-pipeline-num-layers 14 + --decoder-last-pipeline-num-layers 16 + --context-parallel-size 8 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + ... +) +``` + +`TP=4 * PP=8 * CP=8 = 256` 卡构成一个训练组（`DP=1`）。expert group 约束 `expert_tp(1) * EP(32) * PP(8) = 256` 正好整除 world size（`expert_dp=1`）。 + +DSA cross-layer index sharing 要求每个 pipeline stage 都必须**从 computing layer 开始**。在 `index_topk_freq=4` / `index_skip_topk_offset=3` 下，computing layer 是第 1、2、3、7、11、...、75 层。如果直接 `78/8` 均分，stage 会从 skip layer 开始，触发 `get_glm5_spec` 里的 index-share 断言。因此我们使用 `--decoder-first-pipeline-num-layers 14` 和 `--decoder-last-pipeline-num-layers 16`，中间 6 个 stage 各 `(78-14-16)/6 = 8` 层。各 stage 的起始全局层为 1、15、23、31、39、47、55、63，全部是 computing layer。 + +#### BF16 训练 + FP8 Rollout + +训练脚本直接在 `CKPT_ARGS` 和 `ROLLOUT_ARGS` 中写入默认路径，风格与其他示例脚本保持一致： + +```bash +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/GLM-5.2-FP8 + --ref-load $BASE_DIR/GLM-5.2_torch_dist + --load $BASE_DIR/GLM-5.2_vime + --save $BASE_DIR/GLM-5.2_vime + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + ... +) +``` + +`--hf-checkpoint` 提供 vLLM rollout 所需的 FP8 权重和 tokenizer；`--ref-load` 是从 BF16 checkpoint 转换出的 Megatron torch_dist 权重。需要调试 BF16 rollout 时，可以直接把脚本里的 `--hf-checkpoint` 改成 `$BASE_DIR/GLM-5.2`。 + +#### vLLM 配置 + +rollout 侧采用 **prefill/decode (PD) 分离**:1 个 prefill engine(64 卡)+ 3 个 decode engine(192 卡)= 256 卡(必须等于 colocate 的 `rollout_num_gpus`)。每个 engine 使用 64 卡 data parallel 和 vLLM expert parallel。prefill 使用 DeepEP high-throughput backend，decode 使用 low-latency backend。切分通过 `--vllm-config` YAML 配置: + +```yaml +vllm: + - name: default + server_groups: + - worker_type: prefill + num_gpus: 64 + num_gpus_per_engine: 64 + overrides: { data_parallel_size: 64, enable_expert_parallel: true, all2all_backend: deepep_high_throughput, kv_transfer_config: { kv_connector: MooncakeConnector, kv_role: kv_producer, ... }, ... } + - worker_type: decode + num_gpus: 192 + num_gpus_per_engine: 64 + overrides: { data_parallel_size: 64, enable_expert_parallel: true, max_cudagraph_capture_size: 72, all2all_backend: deepep_low_latency, kv_transfer_config: { kv_connector: MooncakeConnector, kv_role: kv_consumer, ... }, ... } +``` + +上游的 `mooncake` 传输对应 vLLM 的 `MooncakeConnector`。上游的 IB device 列表对应 `kv_connector_extra_config.device_name`；prefill 和 decode group 分别使用 `kv_producer` 与 `kv_consumer`。 + +共享 rollout 参数使用 vLLM 原生的 FP8 KV cache 和 CUDA graph 配置: + +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.70 + --vllm-kv-cache-dtype fp8_e4m3 + --vllm-max-cudagraph-capture-size 48 + --vllm-config "${VLLM_CONFIG_FILE}" +) +``` + +MTP / EAGLE speculative decoding 直接使用模型自带的 next-token-prediction 层（GLM-5.2 checkpoint 自带 MTP 层），因此不需要单独的 draft model： + +```bash +--vllm-speculative-config '{"method":"mtp","num_speculative_tokens":5}' +``` + +vLLM 的 CUDA graph capture size 按展开后的 query token 数计算。启用 5 个 speculative token 后，每个 decode 请求对应 `1 + 5 = 6` 个 query token，因此共享上限 `48` 可覆盖 8 个请求，decode group 的覆盖值 `72` 可覆盖 12 个请求。vLLM 会根据 scheduler token capacity 自动推导 DeepEP dispatch buffer 大小。 + +`VLLM_ENGINE_ITERATION_TIMEOUT_S=3600` 会为这个长时间运行的多节点任务提高 vLLM engine watchdog 的超时时间。 + +#### 网络 + +DeepEP/NVSHMEM 的跨节点通信需要在 Ray runtime env 中配置 IB 相关的 NCCL 参数（`NCCL_SOCKET_IFNAME`、`NCCL_IB_*`、`NCCL_NET_GDR_LEVEL`、`NCCL_P2P_LEVEL=NVL`、`NCCL_NVLS_ENABLE=0`、`MC_IB_PCI_RELAXED_ORDERING` 等）。脚本默认使用 `SOCKET_IFNAME=eth0`，如环境不同可在启动前设置 `SOCKET_IFNAME`，它会同时写入 `GLOO_SOCKET_IFNAME`、`TP_SOCKET_IFNAME` 和 `NCCL_SOCKET_IFNAME`。DeepEP 还要求设置 `NVSHMEM_DISABLE_NCCL=1`。 diff --git a/docs/zh/examples/qwen3-30B-A3B.md b/docs/zh/examples/qwen3-30B-A3B.md index 039b3920b..bb85ddc81 100644 --- a/docs/zh/examples/qwen3-30B-A3B.md +++ b/docs/zh/examples/qwen3-30B-A3B.md @@ -30,7 +30,7 @@ bash scripts/run-qwen3-30B-A3B.sh 这里我们简单介绍一下脚本 [run-qwen3-30B-A3B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-30B-A3B.sh) 中与 MoE 相关的部分。 -1. 为了支持在 8xH800 环境中运行 Qwen3-30B-A3B，我们需要开启 megatron 的 CPU Adam 以节省显存，对应配置为： +1. 为了支持在 8xH100 环境中运行 Qwen3-30B-A3B，我们需要开启 megatron 的 CPU Adam 以节省显存，对应配置为： ```bash OPTIMIZER_ARGS=( @@ -68,95 +68,69 @@ bash scripts/run-qwen3-30B-A3B.sh ) ``` - 如果要在 attention 上做 DP 同时在 expert 上做 EP，可以加 `--vllm-data-parallel-size N` - 配合 `--vllm-enable-expert-parallel`。 + 类似地，如果要在 attention 上做 DP 同时在 expert 上做 EP，可以加 + `--vllm-data-parallel-size N` 配合 `--vllm-enable-expert-parallel`。 -### 多机支持 - -以下以 **2台机器、每台8卡（共16GPU）** 为入门示例；脚本与参数可扩展到 **N节点**。多机与单机的主要差异： - -- 训练模型、数据放在所有节点均可访问的路径（如 NFS）； -- `MASTER_ADDR` 设为 head 节点的 **局域网 IP**（非 `127.0.0.1`）； -- 去掉 CPU Adam（多机使用 distributed optimizer，无需 `--optimizer-cpu-offload`）； -- `global-batch-size` 必须等于 `rollout-batch-size × n-samples-per-prompt`。 - -#### 拓扑概览 +### bf16 训练 fp8 推理 -| 组件 | 双机默认配置 | -|------|------| -| 集群 | `ACTOR_NUM_NODES=2`，`ACTOR_NUM_GPUS_PER_NODE=8` | -| Megatron训练 | TP=8, EP=8, CP=2（expert 分片跨节点） | -| vLLM Rollout | 跨节点 TP=16（`rollout-num-gpus-per-engine = 节点数 × 每节点GPU`） | -| 调度 | Ray 集群 + `--colocate` 共卡模式 | +vime 也支持 bf16 训练、fp8 推理。对于 Qwen3-30B-A3B 模型，只需额外下载 fp8 权重： -转换 checkpoint 时建议使用与训练一致的 Megatron 并行度（双机示例 TP=8, EP=8）。checkpoint 的 EP 需与 `--expert-model-parallel-size` 一致，否则 `load_checkpoint` 可能极慢或卡住。 - -#### 启动 Ray 集群 +```bash +hf download Qwen/Qwen3-30B-A3B-FP8 --local-dir /root/Qwen3-30B-A3B-FP8 +``` -Ray 集群需在各节点上 **单独启动**，不在训练脚本内。先在所有 worker 节点加入集群，确认 `ray status` 显示预期 GPU 总数后，再在 head 提交训练。示例（双机）： +并将脚本中的 `--hf-checkpoint` 替换为： ```bash -# === Head 节点 === -export MASTER_ADDR= -ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ - --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# === 各 Worker 节点 === -export MASTER_ADDR= -ray start --address="${MASTER_ADDR}:6379" --node-ip-address=<本机_局域网_IP> --num-gpus 8 +#--hf-checkpoint /root/Qwen3-30B-A3B +--hf-checkpoint /root/Qwen3-30B-A3B-FP8 ``` -更多说明见 [快速开始 — 多机训练](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models)。 +即可触发 fp8 推理。目前我们会将 bf16 权重直接 cast 为 fp8，后续会逐渐加入对精度影响更小的量化方案。 -#### 执行训练 +⚠️ 训练用的 megatron checkpoint 仍需是最初用 bf16 huggingface 权重转换得到的（`--ref-load` / `--load` 不变）。 -Ray 集群就绪后，在 **head 节点** 设置多机环境变量并运行 **与单机相同的脚本**（`ACTOR_NUM_NODES>1` 时脚本不会启动 Ray，并使用多机默认参数）： +### 多机支持 -```bash -export MASTER_ADDR= -export ACTOR_NUM_NODES=2 -export ACTOR_NUM_GPUS_PER_NODE=8 -cd /root/vime -bash scripts/run-qwen3-30B-A3B.sh -``` +以下以 **2 节点 × 8 卡（共 16 GPU）colocate** 为例。多机与单机的差异只在于"跨节点启动 Ray"和"调整几个资源/并行度参数"，训练脚本主体不变。 -2 step 冒烟示例： +1. **共享存储**：模型、数据、checkpoint 放在所有节点路径一致且都能访问的位置（如 NFS）。 -```bash -NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh -``` - -扩展到 N 节点（例如 4×8）时，在各 worker 加入 Ray 后，于 head 设置 `ACTOR_NUM_NODES=4` 并按总卡数调整 `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE`。 +2. **跨节点启动 Ray**（在训练脚本之外，各节点手动执行；详见 [快速开始 — 多机训练](../get_started/quick_start.md#大规模-moe-模型的多机训练)）： -#### 多机关键参数 + ```bash + # Head 节点（node0）；MASTER_ADDR 用局域网 IP，不能是 127.0.0.1 + export MASTER_ADDR= + ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats -| 变量 | 双机默认 | 说明 | -|------|----------|------| -| `ACTOR_NUM_NODES` | 2（单机默认为 1） | 训练节点总数（含 head）；>1 时脚本不启动 Ray | -| `ACTOR_NUM_GPUS_PER_NODE` | 8 | 每节点 GPU 数 | -| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron 并行 | -| `ROLLOUT_NUM_GPUS_PER_ENGINE` | 总 GPU 数 | vLLM engine 占用卡数 | -| `ENABLE_R3` | 1 | 设为 0 可关闭 R3 路径 | + # 其余各节点 + ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 + ``` -脚本默认 batch：`rollout-batch-size=4`，`n-samples-per-prompt=2`，`global-batch-size=8`；vLLM 使用 `--vllm-moe-backend triton`。 + 等 `ray status` 显示 16 GPU 后再提交。由于集群是你手动起好的，运行前需让脚本跳过它开头的进程管理逻辑——把开头的**清理块**（`ray stop --force`、`pkill -9 ray`、`pkill -9 python`、`pkill -9 redis`）**和** `ray start --head ...` 一行都删掉（或注释掉）。否则脚本会先把你刚起的 head 杀掉（并让 worker 变孤儿），导致 `ray job submit` 连 `http://127.0.0.1:8265` 失败。脚本其余部分保留——它仍会 source 模型参数并执行 `ray job submit`。 -#### 多机常见问题 +3. **调整脚本参数**（`scripts/run-qwen3-30B-A3B.sh`）： + - 把 `train.py` 的 `--actor-num-nodes` 由 `1` 改为 `2`（`--actor-num-gpus-per-node` 保持 8）。colocate 下 `--rollout-num-gpus` 会自动取 `actor_num_gpus_per_node × actor_num_nodes = 16`，无需手设。 + - 卡数翻倍后相应增大 `PERF_ARGS` 的并行度（如提高 TP 或引入 DP）；大规模的具体配比可参考 [GLM-4.7](glm4.7-355B-A32B.md)、[DeepSeek-R1](deepseek-r1.md) 等更大集群的例子。 + - `global-batch-size` 必须等于 `rollout-batch-size × n-samples-per-prompt`。 + - （可选）多机使用 distributed optimizer，optimizer 显存压力下降，可去掉 `OPTIMIZER_ARGS` 中的 CPU Adam（`--optimizer-cpu-offload` 等）以提速。 -- **Worker 无法加入 Ray / NCCL 失败**：检查 `MASTER_ADDR`、容器 `/etc/hosts`（hostname 勿指向 `127.0.0.1`）、`NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`。 -- **`Not enough samples X for global_batch_size Y`**：同步调整 `global-batch-size` 与 `rollout-batch-size × n-samples-per-prompt`。 -- **GPU 显存占满但无进程**：重启容器或 `ray stop --force` 清理残留 vLLM 上下文。 +4. **让每个 vLLM engine 留在单节点内**：推荐 `--rollout-num-gpus-per-engine 8`（每节点 1 个 engine），而不是 `16`（单个 engine 跨 2 节点 TP=16）。跨节点 TP 会明显变慢、且对 per-token 数值更敏感；该值需整除总推理卡数（此例为 16）。 -#### EPLB +⚠️ 常见问题： +- **Worker 加不进 Ray / NCCL 失败**：检查 `MASTER_ADDR`、容器 `/etc/hosts`（hostname 勿指向 `127.0.0.1`）、多网卡时设 `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`。 +- **`Not enough samples X for global_batch_size Y`**：保持 `global-batch-size = rollout-batch-size × n-samples-per-prompt`。 +- **每节点少于 8 卡（colocate）**：需显式设置 `--num-gpus-per-node`。 -当总卡数并不能被 expert 总数整除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert。例如对于 24 卡的场景： +此外，当总卡数不能被 expert 总数整除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert。例如 24 卡的场景： - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 24 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 3 - --vllm-enable-expert-parallel - --vllm-enable-eplb - --vllm-eplb-config '{"num_redundant_experts": 16}' - ) - ``` +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 24 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 3 + --vllm-enable-expert-parallel + --vllm-enable-eplb + --vllm-eplb-config '{"num_redundant_experts": 16}' +) +``` diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index e0fac9059..1bf93bc4c 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -2,7 +2,7 @@ ## 环境准备 -拉取 `inferactinc/public:vime-latest` 镜像后，用如下方式初始化镜像环境： +拉取 `vllm/vime:latest` 镜像后，用如下方式初始化镜像环境： ```bash cd /root/ @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM_ARGS -vLLM 推理所需的参数。vime 默认使用 vLLM 作为 rollout 后端（`rollout.py` 启动 `VLLMEngine`，默认 rollout 函数为 `vime.rollout.vllm_rollout.generate_rollout`），无需额外指定 backend。`--rollout-num-gpus-per-engine` 对应每个 vLLM engine 的 `tensor_parallel_size`；除此之外的 vLLM 参数均通过添加 `--vllm-` 前缀传给 vime（例如 `--vllm-max-model-len`）。 +vLLM 推理所需的参数。在默认并行配置下，`--rollout-num-gpus-per-engine` 对应 vLLM 的 `tensor_parallel_size`；除此之外的 vLLM 参数均通过添加 `--vllm-` 前缀传给 vime。 ```bash VLLM_ARGS=( @@ -202,9 +202,7 @@ VLLM_ARGS=( ) ``` -rollout 并发较高时，还可以通过 `--vllm-` 前缀调节 vLLM scheduler，例如 `--vllm-max-num-seqs`、`--vllm-max-num-batched-tokens`；调试或规避 CUDA graph 相关限制时可加 `--vllm-enforce-eager`。 - -⚠️ vime 会用 vLLM router 来调度多个 vLLM server。训推一体（`--colocate`）时，训练与推理权重经 CUDA IPC 同步；训推分离时，训练侧经 NCCL 与 vLLM engine 同步权重。 +⚠️ vime 会用 vllm-router 来调度多个 vLLM server。 ### dynamic sampling @@ -265,7 +263,7 @@ ray job submit ... \ ... ``` -即开启训推一体（colocate），并且训练部分会使用 1 机 8 卡，推理会和训练共同使用这 8 张卡张卡。 +即开启训推一体（colocate），并且训练部分会使用 1 机 8 卡，推理会和训练共同使用这 8 张卡。 如果想使用训推分离的功能，需要去掉 `--colocate` 并配置上 `--rollout-num-gpus`，例如： @@ -278,25 +276,26 @@ ray job submit ... \ ... ``` -此时，就会分配 2 张卡给训练，6 张卡给推理。`--rollout-num-gpus` 与 `--actor-num-gpus-per-node` 一样，是传给 `train.py` 的 **Ray 资源参数**：框架据此创建 placement group，并把前若干 bundle 分给训练 actor、后续 bundle 分给 rollout engine（见 `vime/ray/placement_group.py`）。**共卡模式（`--colocate`）下该参数会被忽略**，并自动设为 `actor_num_gpus_per_node * actor_num_nodes`。请勿把 `--rollout-num-gpus` 写在 `VLLM_ARGS` 中。 +此时，就会分配 2 张卡给训练，6 张卡给推理。 -训推分离时，`VLLM_ARGS` 仅需配置推理后端相关参数，例如： +⚠️ 训推分离时，如果每个 vLLM server 上的并发度太大，超过已配置的 CUDA graph capture size，会影响推理速度。可以用以下 2 种方式进行调整： -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.9 - --vllm-max-num-seqs 256 - --vllm-max-num-batched-tokens 8192 -) -``` +1. 通过 `--vllm-server-concurrency` 限制发给一个 vLLM server 的最大并发量，例如： -如需调试或规避 CUDA graph 相关限制，可额外加上 `--vllm-enforce-eager`。 + ```bash + --vllm-server-concurrency 160 + ``` -⚠️ 在训推一体的训练时，megatron 始终会占据一些显存，需要通过 `--vllm-gpu-memory-utilization` 来降低 vLLM 占据的显存比例，并配合 `--train-memory-margin-bytes` 为训练侧预留空间。 +2. 使用 `--vllm-cudagraph-capture-sizes` 配置 vLLM 初始化的 CUDA graph size，例如： + + ```bash + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + ``` ### 异步训练 当进行训推分离时，你会发现训练和推理的 GPU 总是相互等待着，为了避免这种资源空闲，我们可以开启异步训练。开启的方式即为将启动脚本中的 `train.py` 改变为 `train_async.py`。这样 vime 就会在进行当前 rollout 的训练时进行下一个 rollout 的数据生成了。 -`train.py` 和 `train_async.py` 的差别只在于 train loop 的同步逻辑，我们通过 ray 的异步（`.remote`, `ray.get`）实现了这点。 \ No newline at end of file +⚠️ 异步训练时，vLLM 的性能统计日志与训练日志可能混在一起。可通过 `--vllm-disable-log-stats` 关闭性能统计，并用 `--vllm-uvicorn-log-level warning` 降低服务日志量。 + +`train.py` 和 `train_async.py` 的差别只在于 train loop 的同步逻辑，我们通过 ray 的异步（`.remote`, `ray.get`）实现了这点。 diff --git a/docs/zh/examples/qwen3-4b-base-openhermes.md b/docs/zh/examples/qwen3-4b-base-openhermes.md new file mode 100644 index 000000000..b0bf37d18 --- /dev/null +++ b/docs/zh/examples/qwen3-4b-base-openhermes.md @@ -0,0 +1,85 @@ +# SFT Qwen3-4B-Base + +## 环境准备 + +首先需要我们仿照 [示例：Qwen3-4B 模型](qwen3-4B.md) 创建镜像环境与转换 `Qwen3-4B-Base` 模型。 + +之后，我们处理 sft 数据。这里我们以经典的 [OpenHermes-2.5](https://huggingface.co/datasets/teknium/OpenHermes-2.5) 为例，首先把数据处理成适合 vime 加载的格式，可以用如下的脚本进行处理，增加一个符合 openai message 格式的列，并保存在 `/root/openhermes2_5.parquet`。 + +```python +from datasets import load_dataset + +ds = load_dataset("teknium/OpenHermes-2.5")["train"] + +def convert(sample): + conversations = sample["conversations"] + + def convert_role(role): + if role == "human": + return "user" + elif role == "gpt": + return "assistant" + elif role == "system": + return "system" + else: + raise ValueError(f"Unknown role: {role}") + + messages = [ + { + "role": convert_role(turn["from"]), + "content": turn["value"], + } + for turn in conversations + ] + + return {"messages": messages} + +ds = ds.map(convert) +ds.to_parquet("/root/openhermes2_5.parquet") +``` + +## 执行训练 + +执行训练： + +```bash +cd /root/vime +bash scripts/run-qwen3-4B-base-sft.sh +``` + +### 参数简介 + +可以将 [run-qwen3-4B-base-sft.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-4B-base-sft.sh) 与 [run-qwen3-4B.sh](https://github.com/vllm-project/vime/blob/main/scripts/run-qwen3-4B.sh) 进行对比。会发现除了我们将模型由 instruct 模型换为了 base 模型之外，主要进行了如下的几个调整： + +1. 移除了 `VLLM_ARGS` 和 `GRPO_ARGS`。这是因为 sft 的过程中不需要启动 vllm 或者做 grpo 相关的配置； + +2. 将 `ROLLOUT_ARGS` 改名为了 `SFT_ARGS`，并配置为： + + ```bash + SFT_ARGS=( + --rollout-function-path vime.rollout.sft_rollout.generate_rollout + --prompt-data /root/openhermes2_5.parquet + --input-key messages + --rollout-shuffle + --num-epoch 3 + --rollout-batch-size 128 + --global-batch-size 128 + + --loss-type sft_loss + --calculate-per-token-loss + --disable-compute-advantages-and-returns + --debug-train-only + ) + ``` + + vime 中的 sft 实际上是复用了 vime 的 custom rollout 功能，通过 `--rollout-function-path` 将数据生成部分从使用 vllm 的 RL rollout，切换成了从文件中读取数据的 sft 版本，即 `vime.rollout.sft_rollout.generate_rollout`。 + + 对于 sft 来说，建议将 `rollout_batch_size` 与 `global_batch_size` 设置成相同的，并不要配置 `n_samples_per_prompt`，这样相当于是读一个 batch 就训一个 batch。 + + vime 还支持不同的 loss 类型，我们就是通过 `--loss-type sft_loss` 配置上 sft loss 的。 + + 至于 `--calculate-per-token-loss`，这是因为 vime 默认是以 GRPO 的 per sample mean 进行计算的，而一般 sft 训练都是按一个 batch 的所有不被 mask 的 token 取平均，所以建议配置上。 + + 最后 `--disable-compute-advantages-and-returns` 表示 sft 的过程中不需要预先计算 log prob，`--debug-train-only` 表示不需要初始化 vllm。 + +3. 使用了 `train_async.py` 而不是 `train.py`。这是为了利用异步训练的流程，来实现数据 prefetch。 diff --git a/docs/zh/examples/qwen3-next-80B-A3B.md b/docs/zh/examples/qwen3-next-80B-A3B.md new file mode 100644 index 000000000..127bea966 --- /dev/null +++ b/docs/zh/examples/qwen3-next-80B-A3B.md @@ -0,0 +1,98 @@ +# Qwen3-Next-80B-A3B 训练示例 + +## 环境准备 + +搭建环境、下载模型、数据与 ckpt 转换均与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](./qwen3-4B.md)，将文中 Qwen3-4B 的部分转换为 +Qwen3-next-80B-A3B-Instruct 即可。 + +可以用如下完整方法把 huggingface checkpoint 转化为 torch_dist 格式： + +```bash +export BASE_FOLDER=./models/ +# 下载模型权重 (Qwen3-Next-80B-A3B-Thinking) +hf download Qwen/Qwen3-Next-80B-A3B-Thinking --local-dir ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking +``` + +```shell +cd vime/ +pip install -e . --no-deps + +# (for acceleration) +cd .. # and find a proper folder +git clone https://github.com/fla-org/flash-linear-attention +cd flash-linear-attention +git checkout 9714c595 +pip install -e . --no-deps + +wget https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.4/causal_conv1d-1.5.4+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl +pip install ./causal_conv1d-1.5.4+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl +``` + +## [Optional] Fix a bug in triton compilation on Blackwell (sm100) + +see discussion here https://github.com/triton-lang/triton/issues/8695 +and https://github.com/fla-org/flash-linear-attention/issues/638 + +We need to apply a patch to fix the bug. +Go to the flash-linear-attention folder you just installed, and apply the following patch: + +```diff +diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py +index c5119dcf..838f5e4e 100644 +--- a/fla/ops/gated_delta_rule/wy_fast.py ++++ b/fla/ops/gated_delta_rule/wy_fast.py +@@ -198,7 +198,14 @@ def prepare_wy_repr_bwd_kernel( + b_A += tl.dot(b_kb, tl.trans(b_k)) + b_dkb = tl.dot(b_dA, b_k) + b_db += tl.sum(b_dkb * b_k, 1) +- b_dk += tl.dot(tl.trans(b_dA), b_kb) ++ b_dk += tl.inline_asm_elementwise( ++ asm="mov.f32 $0, $1;", ++ constraints="=r,r", ++ args=[tl.dot(tl.trans(b_dA), b_kb)], ++ dtype=tl.float32, ++ is_pure=True, ++ pack=1, ++ ) + b_dk += b_dkb * b_b[:, None] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + +``` + +save it as `patch.diff` (Please remember to copy the last empty line to the file!) and do `git apply patch.diff` + +## 执行训练 (Megatron) + +**当前暂不支持Blackwell** + +转换模型权重： + +```bash +source scripts/models/qwen3-next-80B-A3B.sh +PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ + tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-Next-80B-A3B-Thinking/ \ + --save /root/Qwen3-Next-80B-A3B-Thinking_torch_dist/ +``` + +单机8卡 + +```bash +cd /root/vime +export BASE_FOLDER=/root +export MASTER_ADDR=127.0.0.1 +ACTOR_NUM_NODES=1 CP_SIZE=1 bash scripts/run-qwen3-next-80B-A3B.sh +``` +如果显存不够，考虑disable `--accumulate-allreduce-grads-in-fp32`，enable `--grad-reduce-in-bf16` + +多机（4x8） + +```bash +cd /root/vime +export BASE_FOLDER=/root +export MASTER_ADDR=your_master_addr +export HOSTFILE=/path/to/hostfile +bash scripts/run-qwen3-next-80B-A3B.sh +``` diff --git a/docs/zh/get_started/agent.md b/docs/zh/get_started/agent.md new file mode 100644 index 000000000..e5510d534 --- /dev/null +++ b/docs/zh/get_started/agent.md @@ -0,0 +1,77 @@ +# Agentic RL 训练路线图 + +vime 的核心定位并不只是跑单轮 RL，而是把高性能训练、vLLM rollout serving、以及可插拔的数据生成接口组合起来，支持 agent 时代常见的多轮工具调用、sandbox 交互、subagent 分支、context compact 和 test-based reward。 + +这篇文档是一个导航页：当你要把 agent workflow 接进 vime 时，先用它判断该看哪些文档和例子。 + +## 从哪里开始 + +| 目标 | 推荐入口 | +| :--- | :--- | +| 给每条 sample 跑自定义 agent loop、tool call、RAG、browser/terminal/sandbox 交互 | [`--custom-generate-function-path`](customization.md#2-自定义生成函数---custom-generate-function-path)、[编写自定义生成函数](quick_start.md#编写自定义生成函数) | +| 做 verifier reward、test-based reward、环境成功判定或外部 reward 服务 | [`--custom-rm-path`](customization.md#3-奖励模型---custom-rm-path)、[编写自定义奖励函数](quick_start.md#编写自定义奖励函数) | +| 一个 prompt 会产生多个训练样本，例如 subagent、multi-agent、context compact | [custom generate 的 fan-out 返回](customization.md#一个-prompt-产生多个训练样本)、[`examples/multi_agent`](../_examples_synced/multi_agent/README.md) | +| agent rollout 有长尾耗时，希望训练不要被最慢样本卡住 | [`examples/fully_async`](../_examples_synced/fully_async/README.md) | +| agent 需要 sandbox、真实代码修改、测试验证和完整端到端样例 | [`examples/coding_agent_rl`](../_examples_synced/coding_agent_rl/README.md) | +| 多轮 agent 需要更高 vLLM serving 吞吐 | [PD 分离](../advanced/pd-disaggregation.md)、[vLLM Config](../advanced/vllm-config.md) | +| 想开启 vLLM 的优化 flag、router 策略或多模型 serving | [vllm 使用方法](usage.md#vllm-使用方法)、[vLLM Config](../advanced/vllm-config.md)、[投机采样](../advanced/speculative-decoding.md)、[低精度训练](../advanced/low-precision.md) | + +## 推荐接入方式 + +大多数 agentic RL 任务应该先从 `--custom-generate-function-path` 开始。这个函数负责把一次 agent 运行转换成 vime 可训练的 `Sample`：填好 `tokens`、`response_length`、`loss_mask`、`status`，并在需要时填好 `reward` 或交给 `--custom-rm-path` 计算。 + +agent workflow 本身可以使用字符串、chat messages、tool calls、环境 observation 或框架自己的事件格式。但训练目标仍然应该是 token based：尽量保留模型实际采样得到的 token ids，并用 `loss_mask` 区分可训练的模型输出和 prompt、template、tool observation、环境文本。 + +如果一次 prompt rollout 只对应一个训练样本，返回一个 `Sample` 即可。如果一次 rollout 会拆成多个训练片段，例如 subagent 轨迹、main-agent 轨迹、compact 前后的片段，则返回 `list[Sample]`，并给这些 sibling samples 设置相同的 `rollout_id`。这样 vime 会在训练 step 切分和 loss 聚合时把它们视作同一次 rollout，而不是重复计数。 + +只有当你需要替换整个 rollout 编排时，才优先考虑 `--rollout-function-path`。典型场景包括：自定义数据源调度、跨 rollout 的后台队列、完全异步生成，或者默认 `vllm_rollout` 的 prompt × sample 结构已经无法表达你的 workflow。 + +## Agent Runtime Adapters + +vime 提供已有 agent runtime 可用的协议 adapter： + +- `vime.agent.adapters.AnthropicAdapter`：Anthropic Messages API，用于 Claude Code 风格 agent。 +- `vime.agent.adapters.OpenAIAdapter`：OpenAI Chat Completions 和 Responses API，用于 OpenAI SDK / OpenAI Agents SDK 风格 client。 + +adapter 是一个便利层，不是单独的 agent framework。它的 contract 是 message history in，sampled tokens out：adapter 会渲染 chat template，用 `input_ids` 和 `return_logprob=True` 调 vLLM，并把返回的 token ids/logprobs 导出为可训练的 trajectory segments；不会从 response text 重新分词恢复训练目标。 + +在自定义 generate 函数里实例化对应协议的 adapter，用 aiohttp 跑它的 `app`，然后通过 adapter 实例管理每次 rollout： + +```python +from vime.agent.adapters import AnthropicAdapter + +adapter = AnthropicAdapter( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, +) + +adapter.open_session(session_id, sampling_defaults=sampling_params) +# Agent client 向 adapter.app 发送请求。 +segments = await adapter.finish_session(session_id) +``` + +多轮 agent 应使用稳定的 `session_id`。adapter 会把它作为 `X-SMG-Routing-Key` 传给 vLLM，让同一个 session 尽量落到同一个 worker，复用 prefix cache。 + +## 沙盒后端 + +Coding-agent 示例内置面向 E2B 兼容服务的 `vime.agent.sandbox.E2BSandbox`。其余 agent 生命周期只依赖 provider-neutral 的 `vime.agent.sandbox.Sandbox` 协议，因此 Docker、Modal 或本地虚拟机后端可以实现同一协议，而无需修改 rollout 逻辑。参见[移植到新的沙盒后端](../_examples_synced/coding_agent_rl/README.md#porting-to-a-new-sandbox-backend)。 + +## Agent Serving 与性能配置 + +agentic rollout 往往比普通单轮 generation 更依赖 serving 配置：上下文更长、多轮请求更多、请求时长分布更重尾，并且可能同时需要 actor、reference、reward 或工具侧模型。 + +- 常规 vLLM server 参数通过 `--vllm-*` 传入。例如 `--max-model-len` 在 vime 中写作 `--vllm-max-model-len`，`--gpu-memory-utilization` 写作 `--vllm-gpu-memory-utilization`。 +- router 参数通过 `--router-*` 传入。多轮 agent 如果需要会话亲和，可以设置 `--router-policy consistent_hash`，让同一个 `sample.session_id` 的多轮请求落到同一个 worker，提高 prefix cache 命中率；否则可使用默认的 `cache_aware` 策略。详见 [多轮 Agent 的会话亲和路由](../advanced/vllm-config.md#多轮-agent-的会话亲和路由)。 +- 更复杂的拓扑使用 `--vllm-config`：它可以描述 PD 分离、多模型 serving、异构 server groups，以及每组不同的 vLLM overrides。 +- 多轮或 agentic RL 通常建议评估 PD 分离。prefill 与 decode 的负载形态不同，拆开后更容易分别扩展资源。 +- 对 rollout 吞吐敏感时，可以继续查看 [投机采样](../advanced/speculative-decoding.md) 和 [低精度训练](../advanced/low-precision.md)。 + +## 参考样例 + +完整的 coding-agent 样例见 [`examples/coding_agent_rl`](../_examples_synced/coding_agent_rl/README.md)。它展示了一个比较接近真实 agent RL 的端到端形态：每条 sample 启动独立 sandbox，agent 使用工具修改代码，生成 `git diff`，再在干净 sandbox 里跑测试得到 reward。 + +这个样例也演示了 agent fan-out 的训练方式：middleware 会把 trajectory 切成 `subagent`、`wipe`（compact 前被冻结的链）和 `final` 等片段，`generate()` 返回 `list[Sample]`，并让这些片段共享同一个 `rollout_id`。 + +如果你只需要更轻量的入门例子，可以先看 [`examples/multi_agent`](../_examples_synced/multi_agent/README.md) 的多 agent 模式。 diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md index 3c4be7c71..2c5ad8a36 100644 --- a/docs/zh/get_started/customization.md +++ b/docs/zh/get_started/customization.md @@ -8,26 +8,26 @@ vime 通过函数路径参数提供了广泛的自定义能力。这些参数允 | 接口参数 | 用途 | | :--- | :--- | -| [`--rollout-function-path`](#1-rollout-函数---rollout-function-path) | 覆盖整个 rollout 生成逻辑。 | -| [`--custom-generate-function-path`](#2-自定义生成函数---custom-generate-function-path) | 仅覆盖生成步骤（例如用于 RAG 或工具调用）。 | -| [`--custom-rm-path`](#3-奖励模型---custom-rm-path) | 实现自定义奖励计算逻辑。 | -| [`--dynamic-sampling-filter-path`](#4-动态采样过滤器---dynamic-sampling-filter-path) | 在动态采样过程中过滤样本（例如 DAPO）。 | -| [`--buffer-filter-path`](#5-buffer-过滤器---buffer-filter-path) | 在训练前过滤 rollout buffer 中的样本。 | -| [`--rollout-sample-filter-path`](#6-rollout-样本过滤器---rollout-sample-filter-path) | 决定单个样本是否参与损失计算。 | -| [`--rollout-all-samples-process-path`](#7-rollout-全样本处理---rollout-all-samples-process-path) | 在 rollout 后处理所有样本（包括被过滤的样本）。 | -| [`--rollout-data-postprocess-path`](#8-rollout-数据后处理---rollout-data-postprocess-path) | 在计算 log probabilities 后对 rollout 数据进行后处理。 | -| [`--custom-loss-function-path`](#9-自定义损失函数---custom-loss-function-path) | 实现自定义训练损失计算。 | -| [`--custom-tis-function-path`](#10-自定义-tisrs-函数---custom-tis-function-path) | 实现用于离策略（off-policy）校正的自定义重要性采样。 | -| [`--custom-pg-loss-reducer-function-path`](#11-自定义-pg-loss-reducer---custom-pg-loss-reducer-function-path) | 自定义 pg_loss 的归约方式（如 Dr.GRPO）。 | -| [`--custom-reward-post-process-path`](#12-奖励后处理---custom-reward-post-process-path) | 在优势计算前对奖励进行自定义后处理。 | -| [`--custom-convert-samples-to-train-data-path`](#13-样本转训练数据---custom-convert-samples-to-train-data-path) | 覆盖样本到训练数据格式的转换逻辑。 | -| [`--custom-rollout-log-function-path`](#14-日志函数) | 训练 rollout 的自定义日志记录。 | -| [`--custom-eval-rollout-log-function-path`](#14-日志函数) | 评估 rollout 的自定义日志记录。 | -| [`--data-source-path`](#15-数据源---data-source-path) | 覆盖 rollout 提示词的数据源。 | -| [`--eval-function-path`](#16-评估函数---eval-function-path) | 专门为评估覆盖 rollout 函数。 | -| [`--custom-megatron-init-path`](#17-megatron-hook) | Megatron 设置后的自定义初始化。 | -| [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hook) | log probability 计算前的自定义逻辑。 | -| [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hook) | 每个训练步骤前的自定义逻辑。 | +| [`--rollout-function-path`](#rollout-function-path) | 覆盖整个 rollout 生成逻辑。 | +| [`--custom-generate-function-path`](#custom-generate-function-path) | 仅覆盖生成步骤（例如用于 RAG 或工具调用）。 | +| [`--custom-rm-path`](#custom-rm-path) | 实现自定义奖励计算逻辑。 | +| [`--dynamic-sampling-filter-path`](#dynamic-sampling-filter-path) | 在动态采样过程中过滤样本（例如 DAPO）。 | +| [`--buffer-filter-path`](#buffer-filter-path) | 在训练前过滤 rollout buffer 中的样本。 | +| [`--rollout-sample-filter-path`](#rollout-sample-filter-path) | 决定单个样本是否参与损失计算。 | +| [`--rollout-all-samples-process-path`](#rollout-all-samples-process-path) | 在 rollout 后处理所有样本（包括被过滤的样本）。 | +| [`--rollout-data-postprocess-path`](#rollout-data-postprocess-path) | 在计算 log probabilities 后对 rollout 数据进行后处理。 | +| [`--custom-loss-function-path`](#custom-loss-function-path) | 实现自定义训练损失计算。 | +| [`--custom-tis-function-path`](#custom-tis-function-path) | 实现用于离策略（off-policy）校正的自定义重要性采样。 | +| [`--custom-pg-loss-reducer-function-path`](#custom-pg-loss-reducer-function-path) | 自定义 pg_loss 的归约方式（如 Dr.GRPO）。 | +| [`--custom-reward-post-process-path`](#custom-reward-post-process-path) | 在优势计算前对奖励进行自定义后处理。 | +| [`--custom-convert-samples-to-train-data-path`](#custom-convert-samples-to-train-data-path) | 覆盖样本到训练数据格式的转换逻辑。 | +| [`--custom-rollout-log-function-path`](#logging-functions) | 训练 rollout 的自定义日志记录。 | +| [`--custom-eval-rollout-log-function-path`](#logging-functions) | 评估 rollout 的自定义日志记录。 | +| [`--data-source-path`](#data-source-path) | 覆盖 rollout 提示词的数据源。 | +| [`--eval-function-path`](#eval-function-path) | 专门为评估覆盖 rollout 函数。 | +| [`--custom-megatron-init-path`](#megatron-hooks) | Megatron 设置后的自定义初始化。 | +| [`--custom-megatron-before-log-prob-hook-path`](#megatron-hooks) | log probability 计算前的自定义逻辑。 | +| [`--custom-megatron-before-train-step-hook-path`](#megatron-hooks) | 每个训练步骤前的自定义逻辑。 | ## 通过 customization 接口实现 agentic workflow @@ -37,18 +37,18 @@ agentic workflow——multi-turn tool use、sandbox interaction、environment fe | 想做的事 | 应使用的接口 | | :--- | :--- | -| 让每条 sample 跑自定义的 agent loop、tool call、RAG、sandbox 执行、browser/terminal 交互或多轮生成，同时复用 vime 默认 rollout loop | [`--custom-generate-function-path`](#2-自定义生成函数---custom-generate-function-path) | -| 实现 verifier reward、test-based reward、environment 成功判定、rule-based reward 或调用外部 reward 服务 | [`--custom-rm-path`](#3-奖励模型---custom-rm-path) | -| 替换整个 rollout 编排（只在 per-sample 自定义不够用时使用） | [`--rollout-function-path`](#1-rollout-函数---rollout-function-path) | -| 控制任务采样、缓冲、回填，或自定义 prompt / task 数据源 | [`--data-source-path`](#15-数据源---data-source-path) | -| 给 agentic 输出附加自定义 loss mask、metadata，或转换成训练数据 | [`--rollout-data-postprocess-path`](#8-rollout-数据后处理---rollout-data-postprocess-path)、[`--custom-convert-samples-to-train-data-path`](#13-样本转训练数据---custom-convert-samples-to-train-data-path) | -| 调试长耗时的 custom generation、verifier、tool call 或 sandbox 调用 | [`vime.utils.trace_utils`](../developer_guide/trace.md) 中的 trace 工具 | +| 让每条 sample 跑自定义的 agent loop、tool call、RAG、sandbox 执行、browser/terminal 交互或多轮生成，同时复用 vime 默认 rollout loop | [`--custom-generate-function-path`](#custom-generate-function-path) | +| 实现 verifier reward、test-based reward、environment 成功判定、rule-based reward 或调用外部 reward 服务 | [`--custom-rm-path`](#custom-rm-path) | +| 替换整个 rollout 编排（只在 per-sample 自定义不够用时使用） | [`--rollout-function-path`](#rollout-function-path) | +| 控制任务采样、缓冲、回填，或自定义 prompt / task 数据源 | [`--data-source-path`](#data-source-path) | +| 给 agentic 输出附加自定义 loss mask、metadata，或转换成训练数据 | [`--rollout-data-postprocess-path`](#rollout-data-postprocess-path)、[`--custom-convert-samples-to-train-data-path`](#custom-convert-samples-to-train-data-path) | +| 调试长耗时的 custom generation、verifier、tool call 或 sandbox 调用 | [`vime.observability.trace_utils`](../developer_guide/trace.md) 中的 trace 工具 | 这一模式的原生示例：[`examples/multi_agent`](../../../examples/multi_agent/README.md) 中基于 `--rollout-function-path` 的多 agent 模式，以及 [`examples/fully_async`](../../../examples/fully_async/README.md) 中适合 long-tail agentic 场景的 fully-async rollout，两者外层都走 vime 默认的 `vllm_rollout`。 ## 详细接口参考 -### 1. Rollout 函数 (`--rollout-function-path`) +### `--rollout-function-path` **默认值**: `vime.rollout.vllm_rollout.generate_rollout` @@ -64,11 +64,11 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False) -> Rollout - 添加自定义采样策略 - 在生成过程中集成外部工具或 API -**示例**: 参见 [examples/multi_agent/rollout_with_multi_agents.py](../../../examples/multi_agent/rollout_with_multi_agents.py) +**示例**: 参见 [examples/fully_async](../_examples_synced/fully_async/README.md) --- -### 2. 自定义生成函数 (`--custom-generate-function-path`) +### `--custom-generate-function-path` **默认值**: `None`（使用内置生成函数） @@ -76,7 +76,7 @@ def generate_rollout(args, rollout_id, data_source, evaluation=False) -> Rollout **函数签名**: ```python -async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample +async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample | list[Sample] ``` **使用场景**: @@ -84,9 +84,41 @@ async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample - 添加检索增强生成（RAG） - 多轮对话处理 +#### 一个 prompt 产生多个训练样本 + +在 subagent、multi-agent、context compact 等 agentic 场景中，一次 prompt rollout 可能会自然拆成多个可训练片段。例如：主 agent 调用 subagent 后，subagent 的轨迹和主 agent 的后续轨迹都需要参与训练；或者发生 compact 后，compact 前后的上下文被切成多个 segment。 + +这种情况下不需要重写整个 rollout 函数，`custom_generate` 可以直接返回 `list[Sample]`。关键是：这些由同一次 rollout 拆出来的 sibling samples 必须设置相同的 `rollout_id`，这样 vime 会在训练切分和 loss 聚合时把它们视作同一次 rollout，而不是重复计数为多次独立 rollout。 + +```python +import copy + +from vime.utils.types import Sample + + +async def custom_generate(args, sample: Sample, sampling_params: dict) -> list[Sample]: + segments = await run_agent_and_split_segments(args, sample, sampling_params) + rollout_id = sample.rollout_id if sample.rollout_id is not None else sample.index + + samples: list[Sample] = [] + for segment in segments: + s = copy.copy(sample) + s.tokens = segment.tokens + s.response = segment.response + s.response_length = segment.response_length + s.loss_mask = segment.loss_mask + s.reward = segment.reward + s.status = Sample.Status.COMPLETED + s.rollout_id = rollout_id + samples.append(s) + return samples +``` + +如果一个完整 trajectory 只有一个总奖励、但被拆成了 `K` 个训练片段，常见做法是在这些片段之间分配这个奖励（例如每个片段写入 `reward / K`），避免把同一次 rollout 的奖励重复放大。 + --- -### 3. 奖励模型 (`--custom-rm-path`) +### `--custom-rm-path` **默认值**: `None`（基于 `--rm-type` 使用内置奖励模型） @@ -118,7 +150,7 @@ async def batched_custom_rm(args, samples: list[Sample]) -> list[float] --- -### 4. 动态采样过滤器 (`--dynamic-sampling-filter-path`) +### `--dynamic-sampling-filter-path` **默认值**: `None` @@ -146,7 +178,7 @@ class DynamicFilterOutput: --- -### 5. Buffer 过滤器 (`--buffer-filter-path`) +### `--buffer-filter-path` **默认值**: `None` @@ -164,7 +196,7 @@ def buffer_filter(args, rollout_id, buffer: list[list[Sample]], num_samples: int --- -### 6. Rollout 样本过滤器 (`--rollout-sample-filter-path`) +### `--rollout-sample-filter-path` **默认值**: `None` @@ -183,7 +215,7 @@ def filter_function(args, samples: list[Sample]) -> None --- -### 7. Rollout 全样本处理 (`--rollout-all-samples-process-path`) +### `--rollout-all-samples-process-path` **默认值**: `None` @@ -200,7 +232,7 @@ def process_function(args, samples: list[list[Sample]], data_source) -> None --- -### 8. Rollout 数据后处理 (`--rollout-data-postprocess-path`) +### `--rollout-data-postprocess-path` **默认值**: `None` @@ -217,7 +249,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 9. 自定义损失函数 (`--custom-loss-function-path`) +### `--custom-loss-function-path` **默认值**: `None`（需要 `--loss-type custom_loss`） @@ -230,7 +262,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 10. 自定义 TIS/RS 函数 (`--custom-tis-function-path`) +### `--custom-tis-function-path` **默认值**: `None` @@ -244,7 +276,7 @@ def postprocess_function(args, samples: list[list[Sample]]) -> None --- -### 11. 自定义 pg_loss Reducer (`--custom-pg-loss-reducer-function-path`) +### `--custom-pg-loss-reducer-function-path` **默认值**: `None` @@ -266,7 +298,7 @@ def get_pg_loss_reducer( --- -### 12. 奖励后处理 (`--custom-reward-post-process-path`) +### `--custom-reward-post-process-path` **默认值**: `None`（使用默认的 GRPO 归一化） @@ -278,7 +310,7 @@ def get_pg_loss_reducer( --- -### 13. 样本转训练数据 (`--custom-convert-samples-to-train-data-path`) +### `--custom-convert-samples-to-train-data-path` **默认值**: `None`（使用内置转换逻辑） @@ -318,7 +350,7 @@ dict: { --- -### 14. 日志函数 +### Logging functions #### 训练 Rollout 日志 (`--custom-rollout-log-function-path`) @@ -340,7 +372,7 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool --- -### 15. 数据源 (`--data-source-path`) +### `--data-source-path` **默认值**: `vime.rollout.data_source.RolloutDataSourceWithBuffer` @@ -371,7 +403,7 @@ class CustomDataSource(DataSource): --- -### 16. 评估函数 (`--eval-function-path`) +### `--eval-function-path` **默认值**: 与 `--rollout-function-path` 相同 @@ -383,7 +415,7 @@ class CustomDataSource(DataSource): --- -### 17. Megatron Hook +### Megatron hooks #### Megatron 初始化 (`--custom-megatron-init-path`) @@ -423,6 +455,24 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler | `--use-routing-replay` | 训练中前向-反向路由一致性。([arXiv:2507.18071](https://arxiv.org/abs/2507.18071)) | | `--use-rollout-routing-replay` | R3：在训练时重放 rollout 阶段的路由。vime 默认的 `vllm_rollout` 路径支持该功能。([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) | +--- + +### 19. Disk 权重同步 Post-Write Hook（`--custom-update-weight-post-write-path`） + +**签名**： +```python +def hook(args, version_dir: str, rollout_engines) -> None +``` + +**用途**：在 disk 权重同步（`--update-weight-transport disk`，full 或 delta 模式）的文件写完之后、 +engine 读取之前，在每个训练 rank 上调用。用于在非 POSIX 共享文件系统上发布写入——例如 commit +一个对象存储挂载——否则其他 host 无法看到这些文件。hook 会在每个 rank 上被调用，需要自行去重 +（例如每个容器只执行一次）。 + +post-write hook 返回前必须保证完整版本目录对读取端可见。host-local 的完整 checkpoint +复制随后直接使用该目录作为来源。delta 机制见 +[Delta 权重同步](../advanced/delta-weight-sync.md)。 + ## 自定义函数路径的测试 vime 现在也提供了一组 CPU 契约测试，用于校验这些 customization 接口。测试会通过字符串形式的导入路径来动态加载组件，因此既能回归仓库内置 hook，也能验证用户通过和训练时完全相同的 CLI 参数传入的自定义实现。 @@ -448,9 +498,8 @@ python -m pytest \ tests/plugin_contracts/test_plugin_runtime_hook_contracts.py ``` -每个测试文件也支持直接通过 `python tests/plugin_contracts/.py` 执行，这样可以和 `run-ci-changed` 保持兼容。 - -CI 中也提供了独立的 `run-ci-plugin-contracts` label，给 PR 打上该标签后会并行运行上述全部四个契约测试（无需 GPU）。 +每个测试文件也支持直接通过 `python tests/plugin_contracts/.py` 执行。 +Buildkite 会在始终运行的 `plugin-contracts` CPU step 中执行全部四个契约测试。 如果你要验证自己的自定义实现，可以直接设置环境变量，例如 `VIME_CONTRACT_ROLLOUT_FUNCTION_PATH`、`VIME_CONTRACT_CUSTOM_RM_PATH`，也可以在直接运行测试文件时传参，例如： diff --git a/docs/zh/get_started/quick_start.md b/docs/zh/get_started/quick_start.md index 8078ccb0f..e999ccb83 100644 --- a/docs/zh/get_started/quick_start.md +++ b/docs/zh/get_started/quick_start.md @@ -8,18 +8,26 @@ ### 硬件支持说明 -**vime** 支持多种 NVIDIA GPU 硬件平台： +**vime** 支持多种硬件平台。 -- **B200 系列**：完全支持，运行步骤与 H 系列完全相同 +**NVIDIA GPU**： + +为目前稳定支持的硬件，包括： + +- **GB200 / GB300 / B200 / 300 系列**：完全支持，运行步骤与 H 系列完全相同 - **H 系列 (H100/H200)**：官方支持，具有完整的 CI 测试保护，运行稳定可靠 **重要说明**： - 最新的 Docker 镜像对 B 卡和 H 卡通用，无需额外配置 - Megatron 后端在 H 卡上具有 CI 保护，经过充分测试验证，推荐生产环境使用 - B 卡基本功能稳定，可作为开发和测试参考，但暂无 CI 保护 -- 两种硬件平台使用完全相同的安装和启动流程 +- 两种 NVIDIA 硬件平台使用完全相同的安装和启动流程 +- 其它卡（如A100/A800）也可以运行，但暂不进行功能维护 + + +**AMD GPU**： -- 对于不方便使用 docker 的场景，请参考 [build_conda.sh](https://github.com/vllm-project/vime/blob/main/build_conda.sh)。 +请参考 [AMD 使用教程](../../en/platform_support/amd_tutorial.md)。 ### 拉取并启动 Docker 容器 @@ -27,12 +35,12 @@ ```shell # 拉取最新镜像 -docker pull inferactinc/public:vime-latest +docker pull vllm/vime:latest # 启动容器 docker run --rm --gpus all --ipc=host --shm-size=16g \ --ulimit memlock=-1 --ulimit stack=67108864 \ - -it inferactinc/public:vime-latest /bin/bash + -it vllm/vime:latest /bin/bash ``` ### 安装 vime @@ -52,7 +60,7 @@ pip install -e . --no-deps ```bash # 下载模型权重 (Qwen3-4B) -hf download zai-org/Qwen3-4B --local-dir /root/Qwen3-4B +hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B # 下载训练数据集 (dapo-math-17k) hf download --repo-type dataset zhuzilin/dapo-math-17k \ @@ -287,8 +295,8 @@ OPTIMIZER_ARGS=( ### VLLM_ARGS: vLLM 服务参数 这部分参数用于配置 vLLM 推理服务。 -- `--rollout-num-gpus-per-engine`: 等同于 vLLM 的 `tp_size`。 -- 其他 vLLM 参数可以通过添加 `--vllm-` 前缀传递给 vime，vime 会自动透传给 vLLM。例如，要设置 vLLM 的 `--log-level INFO` 参数，只需使用 `--vllm-log-level INFO` 即可。 +- `--rollout-num-gpus-per-engine`：单个 rollout engine 使用的 worker GPU 总数；只有 data parallel 和 pipeline parallel 都为 1 时，它才等于 vLLM 的 `tensor_parallel_size`。 +- 其他 vLLM 参数可以通过添加 `--vllm-` 前缀传递给 vime，vime 会自动透传给 vLLM。例如，要设置 vLLM 的 `--uvicorn-log-level info` 参数，只需使用 `--vllm-uvicorn-log-level info`。 > ⚠️ **注意**： > vime 使用 `vllm-router` 调度多个 vLLM 引擎。`dp_size` 会通过 `rollout-num-gpus / rollout-num-gpus-per-engine` 计算得到。 @@ -303,7 +311,7 @@ VLLM_ARGS=( ### Colocated Actor and Rollout -在默认的配置下，训练（Actor）和推理（Rollout）的资源是分开指定的，通过 ray 给训练部分分配 `actor_num_nodes * actor_num_gpus_per_node` 张 GPU，给推理分配 `rollout_num_gpus` 张 GPU，也即训推分离。 +在默认的配置下，训练（Actor）和推理（Rollout）的资源是分开指定的，通过 ray 给训练部分分配 `actor_num_nodes * actor_num_gpus_per_node` 张 GPU，给推理分配 `rollout_num_gpus` 张 GPU，也即训推分离。将 `--rollout-num-gpus` 显式设置为 `0` 时，vime 仍会解析 vLLM 参数并启动 router，但不会启动本地 vLLM server。 **标准（分离）配置**： ```bash @@ -317,7 +325,7 @@ ray job submit ... \ 上述配置中，Actor 使用 4 张卡，Rollout 也使用 4 张卡，两者并行运行。 **训推一体化（Colocated）配置**： -要将训练和推理部署在同一组 GPU 上，请添加 `--colocate` 参数，开启后会忽略 `--rollout-num-gpus` 让训练和推理的卡数相等。 +要将训练和推理部署在同一组 GPU 上，请添加 `--colocate` 参数，开启后默认会让训练和推理的卡数相等；也可以显式设置一个不同的正数，例如让 rollout 卡数多于 actor，多出的 GPU 会作为 rollout-only 资源使用。如果显式设置 `--rollout-num-gpus 0`，则只启动 router，不启动本地 vLLM server。 ```bash @@ -525,7 +533,7 @@ CUSTOM_ARGS=( ## 大规模 MOE 模型的多机训练 -为了启动多机任务，首先需要启动一个 ray 集群，即在 node 0 运行： +如果使用 Ray 进行多机训练，可以参考下面的方式启动集群： ```bash # Node0（HEAD） @@ -536,7 +544,7 @@ ray start --head --node-ip-address ${MASTER_ADDR} \ ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 ``` -在 ray 集群启动后，可以在 node 0 提交任务，例如： +在 Ray 集群启动后，可以在 node 0 提交任务，例如： ```bash ray job submit --address="http://127.0.0.1:8265" \ @@ -553,3 +561,7 @@ ray job submit --address="http://127.0.0.1:8265" \ vime 针对大规模混合专家（MoE）模型的分布式训练进行了深度优化。我们提供了一些端到端的训练案例以供参考： - [示例：8xH100 训练 Qwen3-30B-A3B](../examples/qwen3-30B-A3B.md) +- [示例：8xH100 训练 GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md) +- [示例：32xH100 训练 GLM-5.2](../examples/glm5.2-744B-A40B.md) +- [示例：64xH100 训练 GLM-4.7](../examples/glm4.7-355B-A32B.md) +- [示例：128xH100 训练 DeepSeek-R1](../examples/deepseek-r1.md) diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md index 6d6c63b66..5968c86c7 100644 --- a/docs/zh/get_started/usage.md +++ b/docs/zh/get_started/usage.md @@ -19,23 +19,22 @@ - `--actor-num-gpus-per-node`：RL 的 actor 训练的每个节点有卡； -- `--rollout-num-gpus`：rollout （inference）一共需要多少卡； +- `--rollout-num-gpus`：rollout （inference）一共需要多少卡。设置为 `0` 时，vime 仍会解析 vLLM 参数并启动 router，但不会启动本地 vLLM server； -- `--rollout-num-gpus-per-engine`：每个 inference engine 有多少卡，这个参数会比较像 vLLM 的 `tp_size`，也就是在进行多机 serving 的时候，这个数值应该是总卡数，例如 2 机 16 卡 serving 一个模型，这里的值应该是 16。 +- `--rollout-num-gpus-per-engine`：单个 inference engine 使用的 worker GPU 总数；只有 data parallel 和 pipeline parallel 都为 1 时，它才等于 vLLM 的 `tensor_parallel_size`。例如用 2 机 16 卡 serving 一个模型时，这里的值应为 16。 在默认的配置下，我们会根据这些参数，通过 ray 给训练部分分配 `actor_num_nodes * actor_num_gpus_per_node` 张 GPU，给推理分配 `rollout_num_gpus` 张 GPU，也就是实现了训推分离。 当需要训推一体的时候，还需要配置上： -- `--colocate`：开启训推一体。开启后会忽略 `--rollout-num-gpus` 让训练和推理的卡数相等。 +- `--colocate`：开启训推一体。开启后默认会让训练和推理的卡数相等；也可以显式设置一个不同的正数，例如让 rollout 卡数多于 actor，多出的 GPU 会作为 rollout-only 资源使用。如果显式设置 `--rollout-num-gpus 0`，则只启动 router，不启动本地 vLLM server。 此外，vime 支持 Prefill 和 Decode 的分离部署 (PD Disaggregation)，可以通过设置 `--prefill-num-servers` 参数来指定用于 Prefill 的服务器数量。 ### 选择训练后端 -vime 支持多种训练后端，可以通过 `--train-backend` 参数进行选择： - -- `megatron`（默认）：使用 Megatron-LM 作为训练后端，支持大规模模型的高效训练。 +vime 当前使用 Megatron-LM 作为训练后端。为了兼容已有脚本，仍然可以显式传入 +`--train-backend megatron`。 ### 加载 megatron @@ -148,13 +147,14 @@ vLLM 的加载非常简单，只需要： - 在第一个训练步之前，vime 会把 megatron 里的参数同步给 vLLM，所以 `--hf-checkpoint` 中不需要有最新的训练参数，在续训的时候也不需要更换 hf ckpt； - vLLM 默认会从 huggingface ckpt 中 `config.json` 读取模型的最大 context length，可以使用 `--vllm-max-model-len` 参数来对这个值进行覆盖，从而支持进行更长的推理； - 在训推一体的训练过程中，虽然 megatron 和 vLLM 会先后 offload，但是还是需要为对方留有一些空间，需要通过减小 `--vllm-gpu-memory-utilization` 来调整 vLLM 的显存占用总量。 -- vime 支持透传 vllm-router 的参数，方式是在原参数名前加上 `router` 前缀。例如，vllm-router 的 `--balance-abs-threshold` 参数需要设置为 `--router-balance-abs-threshold`。由于 vllm-router 默认使用 cache-aware routing，可能会导致请求分配不均衡的问题。可以通过设置 `--router-balance-abs-threshold 0` 来强制均衡分配，但这可能会影响多轮对话场景下 prefix cache 的命中率。 +- vime 支持透传 vllm-router 的参数，方式是在原参数名前加上 `router` 前缀。例如，vllm-router 的 `--balance-abs-threshold` 参数需要设置为 `--router-balance-abs-threshold`。由于 vllm-router 默认使用 cache-aware routing，可能会导致请求分配不均衡。可以通过设置 `--router-balance-abs-threshold 0` 来强制均衡分配，但这可能会影响多轮对话场景下 prefix cache 的命中率。对于需要会话亲和的多轮会话，可以设置 `--router-policy consistent_hash`，并为每个会话发送稳定的 `x-session-id`。 +- 如果 vLLM engine 已经由外部系统预启动，可以通过 `--rollout-external-engine-addrs host1:port host2:port` 连接。此时如果训练器和 engine 无法建立 NCCL 权重同步 group，可以使用 `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`，vime 会写完整 HF checkpoint 并调用 vLLM 的 `update_weights_from_disk` 热加载；大模型或跨集群场景可进一步使用 `--update-weight-mode delta --update-weight-transport disk`。详见 [External Rollout Engines 配置路线图](../advanced/external-rollout-engines.md) 和 [Delta 权重同步](../advanced/delta-weight-sync.md)。 对于一些 vLLM 的自定义以及 vime 引入 vLLM 的原理，请见 vLLM 使用方法一节。 ### 数据格式 -目前 vime 只支持加载 `.jsonl` 格式文件，即文件的每一行都是一个 json，一行数据的样例（展开后）为： +vime 支持加载 `.jsonl` 和 `.parquet` 格式文件；读取 Parquet 需要安装 `pyarrow`。两种格式中的每条记录都应包含 `--input-key` 和 `--label-key` 指定的字段。下面是一条 JSONL 数据展开后的示例： ```json { @@ -180,11 +180,26 @@ vLLM 的加载非常简单，只需要： 请注意，这里的 `step_loss_mask`（默认值为 1）字段为 SFT 阶段提供，若设置为 0，则会将该轮 `loss_mask` 设置为 0；若设置为 1，则使用正常 `loss_mask`。 另外我们还提供了一个 metadata_key，默认为 `"metadata"`，读取后我们会把数据中的 metadata 加载进 vime，可能会对自定义数据生成或者自定义 reward model 有帮助。 +如果同一次训练混合了多个数据 source，可以在 metadata 中写入 `source_name`： + +```json +{ + "prompt": "...", + "label": "...", + "metadata": { + "source_name": "math" + } +} +``` + +推荐把 source 标识放在 `metadata["source_name"]` 中；自定义 data source 如果已经动态设置了 `sample.source`，vime 也会识别。rollout 转换成训练数据时，vime 会为每个样本生成 `source_names` 并传到训练侧。source 的读取优先级为动态 `sample.source`、`metadata["source_name"]`，都不存在时为 `"unknown"`。这可以用于自定义 reward、filter、日志统计，以及后续按 source 路由 OPD teacher 等需要分 source 处理的场景。 + ### RL 训练需要的超参 - `--advantage-estimator`: 当前训练需要的 RL 算法，目前支持： - `grpo`（https://arxiv.org/abs/2402.03300）； - `gspo`（https://arxiv.org/abs/2507.18071）； + - `cispo`（https://arxiv.org/abs/2506.13585）； - `reinforce_plus_plus` 与 `reinforce_plus_plus_baseline`（https://arxiv.org/abs/2501.03262）； - `ppo`（https://arxiv.org/abs/1707.06347）。 - `--calculate-per-token-loss`：vime 中默认的方案是 per sample loss，即 `mean(sum(sample_i) / len(sample_i))`，如果需要计算 per token loss，即 `sum(sum(sample_i)) / sum(len(sample_i))`，可以开启 `--calculate-per-token-loss`； @@ -222,35 +237,17 @@ PPO（Proximal Policy Optimization）是经典的 RL 算法，使用 critic 模 --advantage-estimator ppo ``` -**注意：PPO 的 Critic 和 Actor 是并列申请 GPU 的**，在资源分配时需要考虑这一点。具体来说： - -- Critic 模型会独立占用一组 GPU，与 Actor 的 GPU 资源分开； -- 可以通过 `--critic-num-nodes` 和 `--critic-num-gpus-per-node` 来配置 critic 使用的资源； -- 如果不配置 critic 的资源参数，默认会使用与 actor 相同的资源配置。 +**注意：当前 PPO 下 Critic 和 Actor 共享同一组训练 GPU**，资源分配时不需要为 critic 额外预留一组独立 GPU。具体来说： -集群资源分配示例： - -```bash -# Actor 使用 1 个节点，4 张 GPU ---actor-num-nodes 1 ---actor-num-gpus-per-node 4 - -# Critic 使用 1 个节点，4 张 GPU（与 Actor 并列） ---critic-num-nodes 1 ---critic-num-gpus-per-node 4 - -# Rollout 使用 8 张 GPU ---rollout-num-gpus 8 -``` +- PPO 会创建 actor 和 critic 两套训练进程组，但它们会被放到同一组 train placement group 上； +- critic 的训练规模跟随 actor 配置，当前 actor / critic 的 Megatron 并行拓扑必须保持一致； +- PPO 会强制开启 train 侧 offload，使 actor 和 critic 在同一批 GPU 上轮流唤醒和释放显存； +- 当前没有单独配置 critic 训练资源的 CLI 参数，critic 的节点数和每节点 GPU 数会由 actor 配置派生。 -在上述配置下，总共需要 `4 (actor) + 4 (critic) + 8 (rollout) = 16` 张 GPU。 PPO 相关参数： -- `--critic-load`：critic 模型的 checkpoint 路径； -- `--critic-save`：critic 模型的保存路径； -- `--critic-lr`：critic 模型的学习率； -- `--critic-lr-warmup-iters`：critic 模型的 warmup 步数； +- `--megatron-config-path`：通过 YAML 对 actor / critic 分别覆盖 Megatron 参数，例如为 critic 单独设置 `load`、`save`、`lr` 或 warmup 参数； - `--num-critic-only-steps`：训练开始时只训练 critic 的步数； - `--eps-clip`：PPO clip 范围； - `--value-clip`：value loss 的 clip 范围； @@ -327,7 +324,6 @@ vime 支持不同程度的自定义数据生成（rollout）。 output = await post( f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate", { - "model": args.hf_checkpoint, "token_ids": prompt_token_ids, "sampling_params": {"max_tokens": sampling_params["max_new_tokens"]}, } @@ -410,7 +406,7 @@ vllm: **服务器组功能：** - `worker_type`：`regular`、`prefill`、`decode` 或 `placeholder`（预留 GPU 位置但不创建引擎） - `overrides`：vLLM `EngineArgs` 字段覆盖字典，会叠加在 `--vllm-*` CLI 参数之上 -- `num_gpus_per_engine`：每组的 TP 大小覆盖 +- `num_gpus_per_engine`：每组中单引擎的 worker GPU 总数覆盖 ## megatron 使用方法 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 8211928d4..08f305224 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -12,6 +12,21 @@ vime 构建于 `slime `_ 之上，slime 正是 G - DeepSeek V3 系列 (DeepSeek V3, V3.1, DeepSeek R1)； - Llama 3。 +按使用场景开始 +-------------- + +- 第一次使用 vime：:doc:`get_started/quick_start` +- 配置 training 和 rollout 参数：:doc:`get_started/usage` +- 添加 custom generation、reward 或 rollout function：:doc:`get_started/customization` +- 构建 agentic RL workflow：:doc:`get_started/agent` +- 配置生产级 vLLM rollout topology：:doc:`advanced/vllm-config` +- 接入 external rollout engines：:doc:`advanced/external-rollout-engines` +- 以字节级 delta 同步权重：:doc:`advanced/delta-weight-sync` +- 使用 PD disaggregation：:doc:`advanced/pd-disaggregation` +- 使用 BF16 训练 + FP8 rollout 或 FP8 KV cache：:doc:`advanced/low-precision` +- 了解 CI 和可靠性覆盖：:doc:`developer_guide/ci` +- 调试、trace 和 profiling 长时间任务：:doc:`developer_guide/debug`、:doc:`developer_guide/trace`、:doc:`developer_guide/profiling` + .. toctree:: :maxdepth: 1 :caption: 开始使用 @@ -26,21 +41,29 @@ vime 构建于 `slime `_ 之上，slime 正是 G :caption: Dense examples/qwen3-4B.md + examples/glm4-9B.md .. toctree:: :maxdepth: 1 :caption: MoE examples/qwen3-30B-A3B.md + examples/glm5.2-744B-A40B.md + examples/glm4.7-355B-A32B.md + examples/deepseek-r1.md .. toctree:: :maxdepth: 1 :caption: 高级特性 + advanced/on-policy-distillation.md advanced/speculative-decoding.md advanced/reproducibility.md advanced/fault-tolerance.md + advanced/observability.md advanced/pd-disaggregation.md + advanced/external-rollout-engines.md + advanced/delta-weight-sync.md advanced/vllm-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md @@ -51,6 +74,7 @@ vime 构建于 `slime `_ 之上，slime 正是 G _examples_synced/fully_async/README.md _examples_synced/multi_agent/README.md + _examples_synced/coding_agent_rl/README.md .. toctree:: :maxdepth: 1 @@ -60,3 +84,7 @@ vime 构建于 `slime `_ 之上，slime 正是 G developer_guide/debug.md developer_guide/trace.md developer_guide/profiling.md + +.. toctree:: + :maxdepth: 1 + :caption: 硬件平台 diff --git a/examples/README.md b/examples/README.md index 518eab493..2bbdbe5fa 100644 --- a/examples/README.md +++ b/examples/README.md @@ -4,9 +4,17 @@ These examples provide concrete examples to leverage vime in your own RL workflo ## Directory Structure -- **[coding_agent_rl](./coding_agent_rl)**: End-to-end SWE coding-agent RL: a real coding agent (claude-code) edits code in a per-sample sandbox, and the resulting `git diff` is graded against the dataset's test harness. +- **[coding_agent_rl](./coding_agent_rl)**: End-to-end SWE coding-agent RL — a real coding agent (claude-code / codex) edits code in a per-sample sandbox, and the resulting `git diff` is graded against the dataset's test harness. +- **[eval_multi_task](./eval_multi_task)**: Example for supporting evaluation multiple tasks with different configs. - **[fully_async](./fully_async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[geo3k_vlm](./geo3k_vlm)**: Training VLMs on a single-turn reasoning task using GRPO on the GEO3K dataset. - **[geo3k_vlm_multi_turn](./geo3k_vlm_multi_turn)**: VLM multi-turn training on Geo3k dataset. +- **[low_precision](../scripts/low_precision/)**: Launch recipes for FP8/INT4 training and inference. +- **[mem_agent](./mem_agent)**: MemAgent long-context RL — chunk-wise memory update, HotpotQA GRPO training, and RULER-HQA evaluation. - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `vime`. +- **[on_policy_distillation](./on_policy_distillation)**: On-policy distillation (OPD) with an external vLLM teacher or a Megatron-loaded teacher. +- **[dspark](./dspark)**: DSpark speculative decoding draft model training — colocate and non-colocate modes for accelerating RL rollouts. +- **[delta_weight_sync](./delta_weight_sync)**: Non-colocated weight sync that ships only the changed bytes over a shared filesystem (training/inference disaggregation), reloading via the vanilla `update_weights_from_disk` path. +- **[reproducibility](../docs/en/advanced/reproducibility.md)**: Guide to bitwise experiment reproduction using deterministic modes. +- **[tau-bench](./tau-bench)**: Multi-turn tool-use agent training in tau-bench environments. - **[train_infer_mismatch_helper](./train_infer_mismatch_helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS). diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md index a647de582..fd50152ec 100644 --- a/examples/coding_agent_rl/README.md +++ b/examples/coding_agent_rl/README.md @@ -2,11 +2,12 @@ This directory provides an example of running end-to-end **SWE (Software-Engineering) coding-agent RL** with vime: a real coding agent (claude-code CLI) drives `Read/Edit/Grep/Bash/Agent` tools inside a fresh sandbox per sample, the model produces a `git diff`, and the diff is graded against the dataset's test harness in a second clean sandbox (no test-cheating). -Two example files and one shared adapter implement the loop: +Two example files, the shared harness package, and one shared adapter implement the loop: -- `generate.py` — per-sample `generate()` registered via `--custom-generate-function-path`. Boots the sandbox, runs claude-code, captures the diff, scores it, and emits one or more `Sample`s back to vime. -- `vime.agent.adapters.AnthropicAdapter` — the shared Anthropic Messages adapter. claude-code talks to it as if it were Anthropic; the adapter tokenizes the current message history each turn, records prompt/output token snapshots, preserves model-generated tokens (`loss_mask=1`) only while later prompts stitch onto them, masks template/observation tokens (`0`), and emits **three kinds of segments** per trajectory: `subagent` (completed `Task/Agent` dispatch), `wipe` (chain frozen by auto-compact), `final` (tail of the main chain). -- `sandbox.py` — coding-agent/SWE helpers built on `vime.agent.sandbox`: install bootstraps, spawn claude-code, capture patches, and run the fresh-sandbox evaluator. The shared sandbox contract lives in `vime.agent.sandbox.Sandbox`. +- `generate.py` — per-sample `generate()` registered via `--custom-generate-function-path`. Boots the sandbox, prepares the SWE workspace, runs the coding harness (claude-code), captures the diff, scores it, and emits one or more `Sample`s back to vime. +- `vime.agent.adapters.AnthropicAdapter` — the shared Anthropic Messages adapter. claude-code talks to it as if it were Anthropic; the adapter tokenizes the current message history each turn, records prompt/output token snapshots, preserves model-generated tokens (`loss_mask=1`) only while later prompts stitch onto them, and masks template/observation tokens (`0`). Each turn is routed into a per-session message tree inside `vime.agent.trajectory.TrajectoryManager`; any divergence in the prompt prefix forks a new branch, so sub-agent dispatches and auto-compaction are handled as separate root-to-leaf chains. `get_trajectory` linearizes each leaf chain into one `Sample`. +- `vime.agent.harness` — harness-agnostic coding-agent lifecycle (install CLI, write config, spawn detached, poll done-marker). `BaseHarness` defines the contract; `CLAUDE_CODE` / `CODEX` are the shipped implementations. Adding a harness is one new file. The shared sandbox contract lives in `vime.agent.sandbox.Sandbox`. +- `swe.py` — harness-agnostic SWE task layer built on `vime.agent.sandbox`: `prepare_workspace` (pre_commands + PROBLEM_STATEMENT.md), `git_diff` (patch capture), and `evaluate` (fresh-sandbox grading). `SWE_PROMPT` is the task instruction handed to whichever harness runs. `generate.py` owns one `AnthropicAdapter` instance. For each sample it calls `adapter.open_session(...)` before starting claude-code, serves `adapter.app` as @@ -19,10 +20,10 @@ The vime training stack itself follows the standard setup. On top of that you ne 1. **An E2B-compatible sandbox cluster** (or any provider that speaks the E2B SDK). Configure via `E2B_API_KEY` (e.g. the standard `e2b_xxx` key from https://e2b.dev, or any internal endpoint that accepts the same SDK). The official SDK validates this value locally, so internal gateways that ignore auth still need a syntactically valid `e2b_` + 40 hex-character placeholder. 2. **Host-side tarballs** that get uploaded into each sandbox at boot: - - Node 22 (`node-v22.x-linux-x64.tar.xz`) — exported as `SWE_HOST_NODE_TARBALL`. - - Claude Code CLI npm tarball (`anthropic-ai-claude-code-local-linux-x64.tgz`) — exported as `SWE_HOST_CC_TARBALL`. -3. **A sandbox metadata file** (`SWE_SANDBOX_METADATA_FILE`, or the generic `VIME_AGENT_SANDBOX_METADATA_FILE`) — JSON dict whose keys are passed as routing tags when booting an E2B sandbox. Must contain the image key referenced by `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` (e.g. `image`). -4. **Network reachability**: each sandbox dials back to the vime head node's Anthropic adapter over `http://${VIME_HEAD_HOST}:${SHIM_PORT}`. The head host must be reachable from inside the sandboxes (set `VIME_HEAD_HOST` to a routable IP, not `127.0.0.1`). + - Node 22 (`node-v22.x-linux-x64.tar.xz`) — exported as `VIME_AGENT_NODE_TARBALL`. + - Claude Code CLI npm tarball (`anthropic-ai-claude-code-local-linux-x64.tgz`) — exported as `VIME_AGENT_CC_TARBALL`. +3. **An image routing key** (`VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY`, legacy `SWE_SANDBOX_IMAGE_METADATA_KEY` still accepted) — the metadata key your E2B gateway uses to route a boot to a specific image (e.g. `image`). Each sample's `metadata.image` is passed under this key when booting the sandbox. +4. **Network reachability**: each sandbox dials back to the host's Anthropic adapter over `http://${ADAPTER_PUBLIC_HOST}:${ADAPTER_PORT}`. The adapter host must be reachable from inside the sandboxes (set `ADAPTER_PUBLIC_HOST` to a routable IP, not `127.0.0.1`). ## Dataset Format @@ -33,7 +34,7 @@ Standard vime JSONL with three keys: "prompt": "", "label": "", "metadata": { - "image": "swedev/scaleswe.oh.34:", // sandbox image reference + "image": "your-registry/swe-image:", // sandbox image reference "workdir": "/workspace/", // repo path inside the sandbox "problem_statement": "", // exactly one of the following two graders: @@ -57,14 +58,22 @@ cd vime/ export HF_CHECKPOINT=/path/to/Qwen3.6-35B-A3B export REF_MODEL_PATH=/path/to/Qwen3.6-35B-A3B_torch_dist export PROMPT_DATA=/path/to/swe_train.jsonl -export SANDBOX_METADATA_FILE=/path/to/sandbox_metadata.json -export SWE_HOST_NODE_TARBALL=/path/to/node-v22.20.0-linux-x64.tar.xz -export SWE_HOST_CC_TARBALL=/path/to/anthropic-ai-claude-code-local-linux-x64.tgz +export VIME_AGENT_NODE_TARBALL=/path/to/node-v22.20.0-linux-x64.tar.xz +export VIME_AGENT_CC_TARBALL=/path/to/anthropic-ai-claude-code-local-linux-x64.tgz + +# Sandbox provider: +export E2B_API_KEY=e2b_xxx # real key for e2b.dev; a syntactically + # valid placeholder if your gateway ignores auth +export VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY=image # metadata key your gateway routes images by bash examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh ``` -The launcher brings up Ray across all hosts in `/root/mpi_rack_hostfile`, dumps every rollout to `runs/${EXP_TAG}_${STAMP}/rollout_dumps/`, and tees stdout into `runs/${EXP_TAG}_${STAMP}/run.log`. +The launcher fans Ray out to every worker listed in `$HOSTFILE` (default +`/root/mpi_rack_hostfile`, one worker IP per line, reachable over passwordless +SSH as `root`) — create that file (or point `HOSTFILE` at your own) before +launching. It then dumps every rollout to `runs/${EXP_TAG}_${STAMP}/rollout_dumps/` +and tees stdout into `runs/${EXP_TAG}_${STAMP}/run.log`. ## New Arguments @@ -100,19 +109,25 @@ VLLM_ARGS=( All set in the launcher; tune per cluster. +Env vars split by layer. `VIME_AGENT_*` are the reusable agent library's +contract (read inside `vime/agent/`); `SWE_*` are this SWE example's task knobs; +`ADAPTER_*` are host-side deployment/reply-path addresses read only by +`generate.py`. Keep new vars on the prefix that matches the layer that reads them. + | Variable | Default | Meaning | | --- | --- | --- | -| `VIME_HEAD_HOST` | `${MASTER_ADDR}` | Public IP the sandbox uses to reach the Anthropic adapter. **Must be routable from inside the sandbox.** | -| `SHIM_BIND_HOST` / `SHIM_PORT` | `0.0.0.0` / `18001` | Bind address of the adapter shim on the head node. | +| `ADAPTER_PUBLIC_HOST` | `${MASTER_ADDR}` | Public IP the sandbox uses to reach the Anthropic adapter. **Must be routable from inside the sandbox.** | +| `ADAPTER_BIND_HOST` / `ADAPTER_PORT` | `0.0.0.0` / `18001` | Bind address of the Anthropic adapter on the host. | | `E2B_API_KEY` | — | E2B (or compatible) API key. | -| `SWE_SANDBOX_METADATA_FILE` / `VIME_AGENT_SANDBOX_METADATA_FILE` | — | JSON dict of routing metadata passed at sandbox boot. | -| `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` | — | Which key in the metadata file holds the image reference (e.g. `image`). | -| `SWE_HOST_NODE_TARBALL` | — | Host path to Node 22 tarball uploaded into each sandbox. | -| `SWE_HOST_CC_TARBALL` | — | Host path to the Claude Code CLI npm tarball. | -| `SWE_TIME_BUDGET_SEC` | `1800` | Wallclock budget for one agent run. | +| `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` | — | **Required.** Which metadata key the E2B gateway routes images by (e.g. `image`); each sample's `metadata.image` is passed under it. (Legacy `SWE_SANDBOX_IMAGE_METADATA_KEY` still accepted.) | +| `VIME_AGENT_NODE_TARBALL` | — | Host path to Node 22 tarball uploaded into each sandbox. | +| `VIME_AGENT_CC_TARBALL` | — | Host path to the Claude Code CLI npm tarball. | +| `VIME_AGENT_CC_EXTRA_ARGS` | (see launcher) | Extra flags appended to the `claude` CLI invocation — registers the read-only `investigator` sub-agent, disables `WebFetch`/`WebSearch`, disables slash commands. | +| `VIME_AGENT_CC_EXTRA_ENVS` | unset | JSON object of extra env vars exported into the `claude` process — escape hatch for env-only knobs (`MAX_THINKING_TOKENS`, `BASH_MAX_TIMEOUT_MS`, ...). Merged last, so it can also override the built-in defaults. | +| `SWE_AGENT_TIME_BUDGET_SEC` | `1800` | Wallclock budget for the in-sandbox agent CLI itself (think/edit/run). | | `SWE_EVAL_TIMEOUT_SEC` | `600` | Wallclock cap on the evaluator sandbox. | -| `SWE_BOOT_CONCURRENCY` | `6` | Cap on simultaneous sandbox boots (eases h2/SSL long-tail). | -| `SWE_CLAUDE_EXTRA_ARGS` | (see launcher) | Extra flags appended to the `claude` CLI invocation — registers the read-only `investigator` sub-agent, disables `WebFetch`/`WebSearch`, disables slash commands. | +| `SWE_ROLLOUT_GUARD_SEC` | `agent+eval+180` | Outer safety net wrapping the whole rollout (boot + workspace + agent + diff + eval). Auto-derived if unset. | +| `SWE_BOOT_CONCURRENCY` | `16` | Cap on simultaneous sandbox boots (eases h2/SSL long-tail). | | `SWE_CC_PROMPT` | unset | Optional override for the user-turn prompt. Setting this to require sub-agent dispatch is the most reliable way to maximize fan-out. | `--rollout-max-response-len` is the per-turn generation cap passed to each @@ -145,8 +160,8 @@ The Anthropic adapter therefore follows a **string in, token out** contract: Multi-turn agents still force the adapter to tokenize later message histories, because tool observations and claude-code's own compacted messages -arrive as strings. `vime.agent.trajectory.merge_turns` stitches those later -prompts against the saved token stream: +arrive as strings. `vime.agent.trajectory.TrajectoryManager` routes +those later prompts against the saved token stream: - New prompt suffixes that are tool/user/environment context are appended with `loss_mask=0`. @@ -160,15 +175,15 @@ That last case is the important correctness guard. A re-tokenization mismatch can make a string-level conversation look continuous while token-level provenance is broken. vime keeps the context needed to continue the agent, but does not backprop through tokens whose sampled origin can no longer be proven. -The unit tests in `tests/test_agent_trajectory.py` cover matched prefixes, -skipped turns, split-output drift, changed token counts, and prompt-base -restarts. +The unit tests in `tests/test_agent/test_trajectory_manager_branching.py` cover matched +prefixes, skipped turns, split-output drift, changed token counts, and +prompt-base restarts. ## Fan-out Semantics -- `generate()` returns `list[Sample]` — one Sample per trajectory **segment** (`subagent` / `wipe` / `final`). -- Per-trajectory reward is split as `reward / K` across segments; `rollout_id` is shared so the per-rollout-mean loss reducer still counts the trajectory once. -- Sub-agent dispatch increases `K` (each completed `Agent` turn block becomes its own segment), so the effective batch after flatten can be much larger than `rollout_batch_size * n_samples_per_prompt`. +- `generate()` returns `list[Sample]` — one Sample per root-to-leaf chain in the per-session message tree. +- Per-trajectory reward is split as `reward / K` across chains; `rollout_id` is shared so the per-rollout-mean loss reducer still counts the trajectory once. +- Sub-agent dispatch and auto-compaction increase `K` (each prompt-prefix divergence forks a new branch), so the effective batch after flatten can be much larger than `rollout_batch_size * n_samples_per_prompt`. ## Porting to a New Sandbox Backend diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 9f0cc599e..abf274ece 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -1,107 +1,158 @@ """Coding-Agent RL: per-sample generate() function for vime. -Wire-up: - --custom-generate-function-path examples.coding_agent_rl.generate.generate -``generate()`` is intentionally a small four-stage orchestrator: - - 1. ``sandbox.run_claude_code`` prepares the agent sandbox and runs claude-code. - 2. ``sandbox.git_diff`` captures the model-produced patch. - 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. - 4. ``_merge_samples`` combines reward + adapter ``TokenSegment``s, - delegating segment-to-``Sample`` fan-out to ``vime.agent.trajectory``. - -All sandbox-side details live in ``sandbox.py``; the LLM plumbing -(Anthropic <-> vLLM /inference/v1/generate, token capture, 3-kind segment split) uses -``vime.agent.adapters.AnthropicAdapter``. - -Dataset row ``metadata`` schema:: - - image: str # sandbox image - workdir: str # repo path inside the sandbox - problem_statement: str # issue body (falls back to sample.prompt) - swepro: dict|None # SWE-bench Pro test harness (preferred) - eval_cmd: str|None # last-resort: shell command (exit 0 = solved) - -Also accepted (sweb-style rows): ``metadata.remote_env_info.f2p_script`` — -a self-contained Python test file ending in ``sys.exit(pytest.main(...))``. -When ``eval_cmd`` is absent, ``_metadata`` wraps this script into a base64 -materialize-and-run shell command so the existing eval path stays unchanged. - -Env knobs (set in run.sh): - - SWE_HOST_NODE_TARBALL host path to a Node 22 tarball (REQUIRED) - SWE_HOST_CC_TARBALL host path to the Claude Code npm tarball (REQUIRED) - SWE_TIME_BUDGET_SEC 1800 per agent run, wallclock - SWE_EVAL_TIMEOUT_SEC 600 per eval test execution - SHIM_BIND_HOST 0.0.0.0 - SHIM_PORT 18001 - VIME_HEAD_HOST public host the sandboxes use to reach the adapter (REQUIRED) +generate() is a four-stage orchestrator: swe.prepare_workspace + harness.run +-> swe.git_diff -> swe.run_evaluation -> adapter.finish_session. The (harness, +adapter) pair is chosen by the SWE_AGENT env var (claude_code | codex); see +_AGENTS below. +Sandbox-side work is split across three layers: the provider-agnostic sandbox +contract (vime.agent.sandbox), the swappable harness lifecycle +(vime.agent.harness), and the SWE task layer (examples.coding_agent_rl.swe -- +dataset parsing, workspace prep, diff, eval). LLM plumbing (Anthropic / OpenAI +<-> vLLM ``/inference/v1/generate``, token capture, segment split) is the +matching vime.agent.adapters adapter. swe.get_metadata documents the dataset row +schema and produces the md dict consumed below. """ from __future__ import annotations import asyncio -import base64 import logging import os +import random import secrets import time import traceback +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from dataclasses import dataclass from typing import Any -from vime.agent.adapters import AnthropicAdapter -from vime.agent.trajectory import TokenSegment, fan_out_sample_segments +from vime.agent.adapters import AnthropicAdapter, OpenAIAdapter +from vime.agent.aiohttp_threaded import FilteredAccessLogger, run_app_in_thread +from vime.agent.harness import ClaudeCodeHarness, CodexHarness +from vime.agent.sandbox import E2BSandbox from vime.utils.misc import SingletonMeta from vime.utils.processing_utils import load_tokenizer from vime.utils.types import Sample -from . import sandbox -from .aiohttp_threaded import run_app_in_thread +from . import swe logger = logging.getLogger(__name__) +logging.getLogger("e2b").setLevel(logging.WARNING) + +_AGENTS = { + "claude_code": (ClaudeCodeHarness, AnthropicAdapter), + "codex": (CodexHarness, OpenAIAdapter), +} +AGENT_NAME = os.environ.get("SWE_AGENT", "claude_code") +if AGENT_NAME not in _AGENTS: + raise ValueError(f"SWE_AGENT={AGENT_NAME!r} not in {sorted(_AGENTS)}") +HARNESS_CLS, ADAPTER_CLS = _AGENTS[AGENT_NAME] + + +@dataclass(frozen=True) +class SweConfig: + eval_protocol: str # eval-path schema/grader (SWE_EVAL_PROTOCOL) + train_protocol: str # train-path schema/grader (SWE_TRAIN_PROTOCOL) + adapter_public_host: str | None + adapter_bind_host: str + adapter_port: int + fork_merge_threshold: int | None + agent_time_budget_sec: int + eval_timeout_sec: int + rollout_guard_sec: int + boot_concurrency: int + boot_retries: int + + @classmethod + def from_env(cls) -> SweConfig: + agent_time_budget = int(os.environ.get("SWE_AGENT_TIME_BUDGET_SEC", "1800")) + eval_timeout = int(os.environ.get("SWE_EVAL_TIMEOUT_SEC", "600")) + guard = int(os.environ.get("SWE_ROLLOUT_GUARD_SEC", "0") or 0) or (agent_time_budget + eval_timeout + 180) + fork = int(v) if (v := os.environ.get("VIME_FORK_MERGE_MAX_RESPONSE_TOKENS")) else None + return cls( + eval_protocol=os.environ.get("SWE_EVAL_PROTOCOL", swe.PROTOCOL_SCALESWE), + train_protocol=os.environ.get("SWE_TRAIN_PROTOCOL", swe.PROTOCOL_SCALESWE), + adapter_public_host=os.environ.get("ADAPTER_PUBLIC_HOST"), + adapter_bind_host=os.environ.get("ADAPTER_BIND_HOST", "0.0.0.0"), + adapter_port=int(os.environ.get("ADAPTER_PORT", "18001")), + fork_merge_threshold=fork, + agent_time_budget_sec=agent_time_budget, + eval_timeout_sec=eval_timeout, + rollout_guard_sec=guard, + boot_concurrency=int(os.environ.get("SWE_BOOT_CONCURRENCY", "16")), + boot_retries=int(os.environ.get("SWE_BOOT_RETRIES", "2")), + ) + + +CONFIG = SweConfig.from_env() + +_BOOT_SEM = asyncio.Semaphore(CONFIG.boot_concurrency) + + +@asynccontextmanager +async def boot_agent_sandbox(image: str, instance_id: str) -> AsyncIterator[E2BSandbox]: + """Boot a fresh E2B sandbox and install the selected harness toolchain. + + Create the sandbox from the dataset image, install Node 22 + the harness CLI + from host tarballs, retry transient boot/install failures, and close the + sandbox when the caller leaves the context. + """ + sb = None + last_err: Exception | None = None + for attempt in range(CONFIG.boot_retries): + cand = E2BSandbox(image) + try: + async with _BOOT_SEM: + await cand.__aenter__() + try: + await HARNESS_CLS().install_cli(cand) + except BaseException: + await cand.__aexit__(None, None, None) + raise + sb = cand + break + except Exception as e: + last_err = e + logger.warning( + "[coding_agent_rl] %s: provision attempt %d/%d failed: %s: %s", + instance_id, + attempt + 1, + CONFIG.boot_retries, + type(e).__name__, + str(e)[:200], + ) + await asyncio.sleep(1 + attempt + random.random()) + if sb is None: + assert last_err is not None + raise last_err + try: + yield sb + finally: + await sb.__aexit__(None, None, None) -SWE_TIME_BUDGET_SEC = int(os.environ.get("SWE_TIME_BUDGET_SEC", "1800")) -SWE_EVAL_TIMEOUT_SEC = int(os.environ.get("SWE_EVAL_TIMEOUT_SEC", "600")) -# Wall-clock guard for the entire generate() call. Defaults to -# SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 (buffer for sandbox boot, -# diff capture, etc). When exceeded, the in-flight sample is aborted with -# reason `wall_clock_timeout` and the rest of the rollout continues -- this -# isolates a single hung trajectory (e.g. stuck in sandbox.evaluate) so it -# does not kill the whole training step. -SWE_GENERATE_GUARD_SEC = int(os.environ.get("SWE_GENERATE_GUARD_SEC", "0") or 0) or ( - SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 -) -SHIM_BIND_HOST = os.environ.get("SHIM_BIND_HOST", "0.0.0.0") -SHIM_PORT = int(os.environ.get("SHIM_PORT", "18001")) - - -# --------------------------------------------------------------------------- -# Singleton: tokenizer + in-process Anthropic adapter + reducer -# --------------------------------------------------------------------------- -class _State(metaclass=SingletonMeta): +class _AdapterService(metaclass=SingletonMeta): def __init__(self, args) -> None: self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) self.max_context_len = int(getattr(args, "rollout_max_context_len", 0) or 0) self.tool_parser = getattr(args, "vllm_tool_call_parser", None) or None self.reasoning_parser = getattr(args, "vllm_reasoning_parser", None) or None vllm_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" - public_host = os.environ.get("VIME_HEAD_HOST") - if not public_host: + if not CONFIG.adapter_public_host: raise RuntimeError( - "VIME_HEAD_HOST is not set. Export it to the host IP that " - "sandboxes can reach for reverse-connection to the Anthropic adapter. " - "Without it the sandbox cannot dial back and the rollout will " - "silently abort." + "ADAPTER_PUBLIC_HOST is not set. Export it to the host IP that " + "sandboxes can reach for reverse-connection to the adapter; " + "without it the sandbox cannot dial back and the rollout aborts." ) - self.adapter = AnthropicAdapter( + self.adapter = ADAPTER_CLS( tokenizer=self.tokenizer, vllm_url=vllm_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, + fork_threshold_tokens=CONFIG.fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, tearing down the in-flight engine ``/inference/v1/generate`` @@ -109,12 +160,15 @@ def __init__(self, args) -> None: # races with the next release_memory_occupation. self.app_handle = run_app_in_thread( self.adapter.app, - host=SHIM_BIND_HOST, - port=SHIM_PORT, + host=CONFIG.adapter_bind_host, + port=CONFIG.adapter_port, thread_name="anthropic-adapter", - runner_kwargs={"handler_cancellation": True}, + runner_kwargs={ + "handler_cancellation": True, + "access_log_class": FilteredAccessLogger, + }, ) - self.adapter_url = f"http://{public_host}:{self.app_handle.port}" + self.adapter_url = f"http://{CONFIG.adapter_public_host}:{self.app_handle.port}" logger.info( "[coding_agent_rl] tokenizer=%s adapter=%s max_context_len=%s tool_parser=%s reasoning_parser=%s", args.hf_checkpoint, @@ -125,164 +179,112 @@ def __init__(self, args) -> None: ) -# --------------------------------------------------------------------------- -# Trajectory -> Sample conversion -# adapter.finish_session() returns TokenSegments. One trajectory yields >=1 -# segments because the agent may compact + reset mid-run; trajectory.py handles -# the mechanical segment -> Sample fan-out. -# --------------------------------------------------------------------------- -@dataclass(frozen=True) -class RewardResult: - reward: float - is_solved: bool - applied_cleanly: bool - +async def generate(args, base_sample: Sample, sampling_params: dict[str, Any], evaluation: bool = False): + """Per-sample agent function with wall-clock guard (see rollout_guard_sec).""" + state = _AdapterService(args) + protocol = CONFIG.eval_protocol if evaluation else CONFIG.train_protocol + md = swe.get_metadata(base_sample, protocol) + instance_id = md["instance_id"] + if not md["image"] or not md["workdir"]: + return _abort_result(base_sample, "missing_image_or_workdir", instance_id) + reason = swe.evaluability_check(md) + if reason: + return _abort_result(base_sample, f"unevaluatable:{reason}", instance_id) -def _start_session( - state: _State, - sample: Sample, - md: dict[str, Any], - sampling_params: dict[str, Any], -) -> str: - # claude-code inside the sandbox dials back to the adapter with this - # session_id (passed as the Bearer token) so its turns are grouped under - # one chain history. Build from (instance_id, index, group_index) when - # possible; fall back to random hex if either index is missing. - if sample.session_id: - session_id = sample.session_id - elif sample.index is not None and sample.group_index is not None: - session_id = f"cagent-{md['instance_id']}-{sample.index}-{sample.group_index}" - else: - session_id = f"cagent-{md['instance_id']}-{secrets.token_hex(8)}" - sample.session_id = session_id + session_id = base_sample.session_id = _session_id(base_sample, instance_id) state.adapter.open_session( session_id, sampling_defaults=sampling_params, max_context_tokens=state.max_context_len, ) - return session_id - - -def _merge_samples( - *, - sample: Sample, - state: _State, - segments: list[TokenSegment], - reward_result: RewardResult, - elapsed_sec: float, - instance_id: str, -): - if not segments: - return _abort_result(sample, "adapter_session_empty") - - trajectory_metadata = { - **(sample.metadata or {}), - "instance_id": instance_id, - "is_solved": reward_result.is_solved, - "applied_cleanly": reward_result.applied_cleanly, - "elapsed_sec": elapsed_sec, - } - - # All K samples share rollout_id so the loss reducer counts this - # trajectory once. - fanned = fan_out_sample_segments( - sample, - segments, - reward_result.reward, - state.tokenizer, - metadata=trajectory_metadata, - ) - if not fanned: - raise ValueError("fan-out produced no samples") - - logger.info( - "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", - instance_id, - reward_result.reward, - reward_result.is_solved, - reward_result.applied_cleanly, - elapsed_sec, - len(fanned), - ) - return fanned - - -# --------------------------------------------------------------------------- -# Main per-sample agent function -# -# The four calls inside the timeout are the high-level rollout recipe: -# run_claude_code -> git_diff -> sandbox.evaluate -> merge_samples. -# --------------------------------------------------------------------------- -async def generate(args, sample: Sample, sampling_params: dict[str, Any]): - """Per-sample agent function with wall-clock guard. See - SWE_GENERATE_GUARD_SEC docstring above.""" - state = _State(args) - md = _metadata(sample) - if not md["image"] or not md["workdir"]: - return _abort_result(sample, "missing_image_or_workdir") - - instance_id = md["instance_id"] - session_id = _start_session(state, sample, md, sampling_params) t0 = time.time() try: - async with asyncio.timeout(SWE_GENERATE_GUARD_SEC): - async with sandbox.boot_agent_sandbox(md["image"]) as sb: - await sandbox.run_claude_code( + async with asyncio.timeout(CONFIG.rollout_guard_sec): + async with boot_agent_sandbox(md["image"], instance_id) as sb: + await swe.prepare_workspace(sb, md["workdir"], md) + agent_exit_code = await HARNESS_CLS().run( sb, workdir=md["workdir"], session_id=session_id, adapter_url=state.adapter_url, - time_budget_sec=SWE_TIME_BUDGET_SEC, - problem_statement=md["problem_statement"], - swepro=md["swepro"], - pre_commands=md["pre_commands"], + time_budget_sec=CONFIG.agent_time_budget_sec, + prompt=swe.SWE_PROMPT, ) - diff_text = await sandbox.git_diff(sb, md["workdir"]) + diff_text = await swe.git_diff(sb, md["workdir"]) - reward, is_solved, applied_cleanly = await sandbox.evaluate( - image=md["image"], - workdir=md["workdir"], + reward, applied_cleanly = await swe.run_evaluation( + md, diff_text=diff_text, - swepro=md["swepro"], - eval_cmd=md["eval_cmd"], - pre_commands=md["pre_commands"], - timeout_sec=SWE_EVAL_TIMEOUT_SEC, + timeout_sec=CONFIG.eval_timeout_sec, ) - reward_result = RewardResult( + if evaluation: + logger.info( + "[coding_agent_rl] %s: reward=%.2f applied=%s agent_exit_code=%d elapsed=%.1fs (eval-only)", + instance_id, + float(reward), + bool(applied_cleanly), + agent_exit_code, + time.time() - t0, + ) + return _eval_result( + base_sample, + reward=float(reward), + applied_cleanly=bool(applied_cleanly), + agent_exit_code=agent_exit_code, + instance_id=instance_id, + ) + + samples = await state.adapter.finish_session( + session_id, + base_sample=base_sample, reward=float(reward), - is_solved=bool(is_solved), - applied_cleanly=bool(applied_cleanly), + extra_metadata={ + "grading_solved": float(reward) == 1.0, + "instance_id": instance_id, + }, ) - segments = await state.adapter.finish_session(session_id) - return _merge_samples( - sample=sample, - state=state, - segments=segments, - reward_result=reward_result, - elapsed_sec=time.time() - t0, - instance_id=instance_id, + if not samples: + return _abort_result(base_sample, "adapter_session_empty", instance_id) + + for s in samples: + s.metadata = {**(s.metadata or {}), "agent_exit_code": agent_exit_code} + if agent_exit_code != 0: + reason = "time budget exceeded" if agent_exit_code < 0 else f"CLI error (exit {agent_exit_code})" + logger.warning( + "[coding_agent_rl] %s: agent_exit_code=%d (%s)", + instance_id, + agent_exit_code, + reason, + ) + logger.info( + "[coding_agent_rl] %s: reward=%.2f applied=%s agent_exit_code=%d elapsed=%.1fs segments=%d", + instance_id, + float(reward), + bool(applied_cleanly), + agent_exit_code, + time.time() - t0, + len(samples), ) + return samples except asyncio.TimeoutError: - _log_timeout_diagnostic(t0) - return _abort_result(sample, "wall_clock_timeout") + _log_timeout_diagnostic(t0, instance_id) + return _abort_result(base_sample, "wall_clock_timeout", instance_id) except Exception as e: - logger.error( + logger.warning( "[coding_agent_rl] %s: rollout failed: %s\n%s", instance_id, e, traceback.format_exc(), ) - return _abort_result(sample, f"exception:{type(e).__name__}") + return _abort_result(base_sample, f"exception:{type(e).__name__}", instance_id) finally: - # Close the sid before next train step's release_memory_occupation; - # stragglers from this trajectory would otherwise race its idle assert. - await state.adapter.finish_session(session_id) # idempotent + await state.adapter.drop_session(session_id, wait_timeout=30) # cleanup only, idempotent + await asyncio.sleep(10) -def _log_timeout_diagnostic(t0: float) -> None: - """Dump pending-task names when the wall-clock guard fires so future - debugging can see which await was stuck. Must never crash.""" +def _log_timeout_diagnostic(t0: float, instance_id: str) -> None: + # Dump pending-task names when the wall-clock guard fires. Must not crash. try: elapsed = time.time() - t0 pending = [t for t in asyncio.all_tasks() if not t.done()] @@ -291,10 +293,11 @@ def _log_timeout_diagnostic(t0: float) -> None: coro = getattr(t, "_coro", None) stuck.append(getattr(coro, "__qualname__", repr(coro))) logger.warning( - "[coding_agent_rl] generate() wall_clock_timeout after %.1fs " + "[coding_agent_rl] %s: wall_clock_timeout after %.1fs " "(guard=%ds); %d tasks pending; sample of stuck: %s", + instance_id, elapsed, - SWE_GENERATE_GUARD_SEC, + CONFIG.rollout_guard_sec, len(pending), stuck, ) @@ -302,62 +305,57 @@ def _log_timeout_diagnostic(t0: float) -> None: pass -# --------------------------------------------------------------------------- -# Metadata helpers -# --------------------------------------------------------------------------- -def _wrap_f2p_script(script: str | None) -> str | None: - # Materialize a self-contained pytest script (typical sweb f2p_script: - # ends with `sys.exit(pytest.main([...]))`) into the sandbox via base64 - # so we sidestep all shell quoting; python's exit code carries the - # pytest pass/fail signal that `_run_eval_cmd` turns into reward. - if not script: - return None - b64 = base64.b64encode(script.encode("utf-8")).decode("ascii") - return f"echo {b64} | base64 -d > /tmp/vime_f2p.py && python /tmp/vime_f2p.py" - - -def _metadata(sample: Sample) -> dict[str, Any]: - """Normalize the two dataset schemas (flat vs ``remote_env_info``).""" - m = sample.metadata or {} - rem = m.get("remote_env_info") or {} - label = sample.label if (isinstance(sample.label, str) and len(sample.label) < 256) else None - return { - "instance_id": m.get("instance_id") or rem.get("instance_id") or label or "unknown", - "image": m.get("image") or rem.get("image_url"), - "workdir": m.get("workdir") or rem.get("workdir"), - "problem_statement": m.get("problem_statement") or _coerce_prompt(sample.prompt), - "swepro": m.get("swepro"), - "eval_cmd": m.get("eval_cmd") or _wrap_f2p_script(rem.get("f2p_script")), - "pre_commands": m.get("pre_commands") or rem.get("pre_commands"), - } - - -def _coerce_prompt(prompt) -> str: - if isinstance(prompt, str): - return prompt - if isinstance(prompt, list): - for m in prompt: - if isinstance(m, dict) and m.get("role") == "user": - c = m.get("content") - if isinstance(c, str): - return c - if isinstance(c, list): - return "\n".join(p.get("text", "") for p in c if isinstance(p, dict) and p.get("type") == "text") - return "" +def _session_id(sample: Sample, instance_id: str) -> str: + if sample.session_id: + return sample.session_id + if sample.index is not None and sample.group_index is not None: + return f"cagent-{instance_id}-{sample.index}-{sample.group_index}" + return f"cagent-{instance_id}-{secrets.token_hex(8)}" -def _abort(sample: Sample, reason: str) -> Sample: +def _abort_result(sample: Sample, reason: str, instance_id: str) -> list[Sample]: + """Mark ``sample`` aborted in place and return it in the list shape this + fan-out generate function always yields.""" sample.tokens = [0, 0] sample.response = "" sample.response_length = 1 sample.loss_mask = [0] + sample.rollout_log_probs = [0.0] sample.reward = 0.0 + sample.remove_sample = True sample.status = Sample.Status.ABORTED - sample.metadata = {**(sample.metadata or {}), "abort_reason": reason} - logger.warning("[coding_agent_rl] aborted: %s", reason) - return sample + sample.metadata = { + **(sample.metadata or {}), + "abort_reason": reason, + "instance_id": instance_id, + } + logger.warning("[coding_agent_rl] %s aborted: %s", instance_id, reason) + return [sample] -def _abort_result(sample: Sample, reason: str): - """Return a uniform list shape for this fan-out generate function.""" - return [_abort(sample, reason)] +def _eval_result( + sample: Sample, + *, + reward: float, + applied_cleanly: bool, + agent_exit_code: int | None, + instance_id: str, +) -> list[Sample]: + """Eval-path placeholder: only ``reward`` matters for ``eval/sweb``.""" + + sample.tokens = [0, 0] + sample.response = "" + sample.response_length = 1 + sample.loss_mask = [0] + sample.rollout_log_probs = [0.0] + sample.reward = float(reward) + sample.remove_sample = True + sample.status = Sample.Status.COMPLETED + sample.metadata = { + **(sample.metadata or {}), + "instance_id": instance_id, + "grading_solved": float(reward) == 1.0, + "applied_cleanly": applied_cleanly, + "agent_exit_code": agent_exit_code, + } + return [sample] diff --git a/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh old mode 100644 new mode 100755 index a307e7b1e..59b1ebdcf --- a/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh +++ b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh @@ -1,33 +1,11 @@ #!/usr/bin/env bash -# End-to-end SWE coding-agent RL on 8 nodes. -# -# Same model and training loop as run_qwen36_35b_a3b_swe_8node.sh, with three -# extra layers that actively encourage the rollout to dispatch sub-agents. -# Trajectory trees produced by this script show real `sibling` branches: -# -# (1) An `investigator` sub-agent is registered via claude-code's --agents -# flag (Grep/Read/Glob only, no edits) — a concrete, narrowly-scoped -# dispatch target. -# (2) SWE_CC_PROMPT requires the model to dispatch the investigator before -# any edit, naming the exact call form (Agent tool with -# subagent_type=investigator). -# (3) Agent/Task tools stay in the allowed set; WebFetch/WebSearch are -# disabled (sandbox has no outbound internet); --disable-slash-commands -# removes /compact as a competing branching pathway. -# -# Fan-out semantics: -# * generate() returns list[Sample] (one Sample per trajectory segment); -# the per-trajectory reward is split as reward/K across segments. -# * Sub-agent dispatch increases K (each sub-agent turn block becomes its -# own segment), so the effective batch after flatten can be much larger -# than rollout_batch_size * n_samples_per_prompt. If pinned-memory or -# GPU wake_up OOM appears, lower rollout_batch_size or n_samples_per_prompt -# first — not max-tokens-per-gpu. -# Run from a long-lived shell / tmux session on the Ray head node; do not wrap -# in a short-lived nohup launcher or Ray child processes get cleaned up with it. +# End-to-end SWE coding-agent RL on 8 nodes. See README.md for the dataset +# schema, env vars, and fan-out semantics. Run from a long-lived shell / tmux +# session on the Ray head node (a short-lived nohup launcher gets its Ray child +# processes cleaned up with it). # Best-effort cleanup so a rerun does not collide with stale workers. -pkill -9 -f "vllm serve" || true +pkill -9 -f '[v]llm serve|VLL[M]::' || true sleep 3 ray stop --force || true pkill -9 ray || true @@ -40,8 +18,6 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" VIME_DIR="${VIME_DIR:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" # ============ model parallelism ============ -# CP=8 (higher than the baseline script's CP=2) gives more rank-local context -# room for the longer per-segment payloads typical under sub-agent dispatch. export TP_SIZE="${TP_SIZE:-2}" export PP_SIZE="${PP_SIZE:-1}" export CP_SIZE="${CP_SIZE:-8}" @@ -72,11 +48,9 @@ MAX_CONTEXT_LEN="${MAX_CONTEXT_LEN:-96000}" MAX_GEN_LEN="${MAX_GEN_LEN:-32768}" # ============ paths — override before launching ============ -# Point these at your own checkpoints / dataset / sandbox metadata. HF_CHECKPOINT="${HF_CHECKPOINT:-/path/to/Qwen3.6-35B-A3B}" REF_MODEL_PATH="${REF_MODEL_PATH:-/path/to/Qwen3.6-35B-A3B_torch_dist}" PROMPT_DATA="${PROMPT_DATA:-/path/to/swe_train.jsonl}" -SANDBOX_METADATA_FILE="${SANDBOX_METADATA_FILE:-/path/to/sandbox_metadata.json}" EXP_TAG="${EXP_TAG:-agent_only}" STAMP="$(date +%Y%m%d_%H%M%S)" @@ -169,21 +143,21 @@ PERF_ARGS=( --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 - # one CP rank's slice of MAX_CONTEXT_LEN; log-probs chunked along T to - # avoid OOM on long single trajectories. + # max-tokens-per-gpu is one CP rank's slice of MAX_CONTEXT_LEN; log-probs are + # chunked along T to avoid OOM on long single trajectories. --max-tokens-per-gpu $((MAX_CONTEXT_LEN / CP_SIZE)) --log-probs-chunk-size 1024 --use-dynamic-batch-size ) ALGO_ARGS=( - --advantage-estimator gspo + --advantage-estimator grpo --kl-loss-coef 0.00 --kl-loss-type low_var_kl --kl-coef 0.00 --entropy-coef 0.00 - --eps-clip 1e-4 - --eps-clip-high 2e-4 + --eps-clip 0.2 + --eps-clip-high 0.28 ) OPTIMIZER_ARGS=( @@ -198,7 +172,6 @@ OPTIMIZER_ARGS=( --use-precision-aware-optimizer ) -# ============ rollout engine ============ VLLM_ARGS=( --rollout-num-gpus 64 --rollout-num-gpus-per-engine ${ROLLOUT_TP_SIZE} @@ -207,8 +180,6 @@ VLLM_ARGS=( --vllm-enable-expert-parallel --vllm-tool-call-parser qwen3_coder --vllm-reasoning-parser qwen3 - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - --prefill-num-servers 1 ) MISC_ARGS=( @@ -223,7 +194,7 @@ MISC_ARGS=( ) # ============ ray cluster network ============ -# Set MASTER_ADDR before the SWE block: VIME_HEAD_HOST below falls back to it. +# Set MASTER_ADDR before the SWE block: ADAPTER_PUBLIC_HOST below falls back to it. export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" @@ -231,63 +202,35 @@ export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" # ============ SWE / claude-code rollout knobs ============ -# --- sandbox provisioning (E2B) --- -# The E2B SDK validates the API key format locally before it reaches -# E2B-compatible gateways. If your internal gateway ignores auth, use any -# syntactically valid e2b_... hex placeholder. -E2B_DUMMY_API_KEY="e2b_0000000000000000000000000000000000000000" -E2B_API_KEY="${E2B_API_KEY:-${E2B_DUMMY_API_KEY}}" -if [[ ! "${E2B_API_KEY}" =~ ^e2b_[0-9a-fA-F]{40}$ ]]; then - echo "WARN: E2B_API_KEY does not pass local E2B SDK format validation; using dummy key." >&2 - E2B_API_KEY="${E2B_DUMMY_API_KEY}" -fi -export E2B_API_KEY -export SWE_SANDBOX_METADATA_FILE="${SANDBOX_METADATA_FILE}" -export SWE_SANDBOX_IMAGE_METADATA_KEY="${SWE_SANDBOX_IMAGE_METADATA_KEY:-glm-platform/image}" -# Host-side tarballs injected into each sandbox at boot. -export SWE_HOST_NODE_TARBALL="${SWE_HOST_NODE_TARBALL:-/path/to/node-v22.x-linux-x64.tar.xz}" -export SWE_HOST_CC_TARBALL="${SWE_HOST_CC_TARBALL:-/path/to/anthropic-ai-claude-code-local-linux-x64.tgz}" - -# --- reply path (sandbox -> host shim) --- -export VIME_HEAD_HOST="${VIME_HEAD_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" -export SHIM_BIND_HOST="${SHIM_BIND_HOST:-0.0.0.0}" -export SHIM_PORT="${SHIM_PORT:-18001}" - -# --- per-trajectory time / concurrency budgets --- -# Time budget 1800s (vs baseline 1200): sub-agent dispatch on large repos blows -# past a tighter budget — investigator passes are the long tail. -# Boot concurrency 6 (vs baseline 8) eases h2/SSL long-tail stalls under -# heavier sub-agent dispatch. -export SWE_TIME_BUDGET_SEC="${SWE_TIME_BUDGET_SEC:-1800}" +export SWE_AGENT="${SWE_AGENT:-claude_code}" +export SWE_TRAIN_PROTOCOL="${SWE_TRAIN_PROTOCOL:-scaleswe}" +export E2B_API_KEY="${E2B_API_KEY:-e2b_0000000000000000000000000000000000000000}" +# Metadata key your gateway routes images by; `image` is the neutral default. +export VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY="${VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY:-image}" +export VIME_AGENT_NODE_TARBALL="${VIME_AGENT_NODE_TARBALL:-/path/to/node-v22.x-linux-x64.tar.xz}" +export VIME_AGENT_CC_TARBALL="${VIME_AGENT_CC_TARBALL:-/path/to/anthropic-ai-claude-code-local-linux-x64.tgz}" + +# ADAPTER_PUBLIC_HOST must be routable from inside the sandbox (not 127.0.0.1). +export ADAPTER_PUBLIC_HOST="${ADAPTER_PUBLIC_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export ADAPTER_BIND_HOST="${ADAPTER_BIND_HOST:-0.0.0.0}" +export ADAPTER_PORT="${ADAPTER_PORT:-18001}" + +export SWE_AGENT_TIME_BUDGET_SEC="${SWE_AGENT_TIME_BUDGET_SEC:-1800}" export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" -export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-6}" - -# --- trajectory fan-out --- -# generate() emits one Sample per segment (reducer splits reward/K); -# rollout_id is shared so the per-rollout-mean loss reducer still counts -# the trajectory once. -# --rollout-max-response-len caps one model turn. The custom generate function -# uses --rollout-max-context-len as the multi-turn prompt+response budget. - -# --- claude-code CLI extras --- -# SETTINGS_JSON: autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI -# compacts before any segment crosses the training-side cap. -# AGENTS_JSON: register a read-only `investigator` sub-agent (Grep/Read/Glob) -# as a concrete, narrowly-scoped dispatch target. -# SWE_CLAUDE_EXTRA_ARGS: WebFetch/WebSearch are off (sandbox has no outbound -# internet); --disable-slash-commands keeps the model from emitting /compact -# as a competing branching pathway. +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-16}" + +# autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI compacts before any +# segment crosses the training-side cap. `investigator` is a read-only sub-agent +# (a concrete dispatch target). WebFetch/WebSearch off (no outbound internet). SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' AGENTS_JSON='{"investigator":{"description":"Searches the repo for relevant files before any edit","prompt":"You are an investigator sub-agent. Use Grep/Read/Glob to find every file relevant to the user task, then return a short bulleted summary. Do NOT edit anything.","tools":["Grep","Read","Glob"]}}' -export SWE_CLAUDE_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" +export VIME_AGENT_CC_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" -# Optional: bias the model to dispatch the investigator before any edit. -# Uncomment to maximize sub-agent dispatch — naming the exact call form -# (Agent tool with subagent_type=investigator) is what reliably triggers it. +# Optional: require dispatching the investigator before any edit, to maximize sub-agent fan-out. # export SWE_CC_PROMPT="Read PROBLEM_STATEMENT.md. BEFORE editing any file, dispatch the 'investigator' sub-agent (via the Agent tool with subagent_type=investigator) to locate every file relevant to the issue. Then fix the issue and run the tests." # ============ proxy bypass for in-cluster traffic ============ -export no_proxy="127.0.0.1,${MASTER_ADDR},${VIME_HEAD_HOST}" +export no_proxy="127.0.0.1,${MASTER_ADDR},${ADAPTER_PUBLIC_HOST}" export NO_PROXY="${no_proxy}" cd "${VIME_DIR}" @@ -306,7 +249,7 @@ if [[ -f "${HOSTFILE}" ]]; then [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue echo "Starting Ray worker on ${WORKER_IP}" ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ - "pkill -9 -f 'vllm serve' ; ray stop --force ; pkill -9 python ; \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; \ ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ --node-ip-address ${WORKER_IP} --disable-usage-stats" & done @@ -323,13 +266,16 @@ RUNTIME_ENV_JSON=$(python3 - < AsyncIterator[E2BSandbox]: - """Boot a fresh E2B sandbox and install the Claude Code toolchain. - - This is the provisioning wrapper for the work sandbox: create the sandbox - from the dataset image, install Node 22 + Claude Code CLI from host - tarballs, retry transient boot/install failures, and close the sandbox when - the caller leaves the context. - """ - global _BOOT_SEM - if _BOOT_SEM is None: - _BOOT_SEM = asyncio.Semaphore(SWE_BOOT_CONCURRENCY) - - sb = None - last_err: Exception | None = None - for attempt in range(SWE_BOOT_RETRIES): - cand = E2BSandbox(image) - try: - async with _BOOT_SEM: - await cand.__aenter__() - try: - await install_node22(cand, SWE_HOST_NODE_TARBALL) - await install_claude_code(cand, SWE_HOST_CC_TARBALL) - except BaseException: - await cand.__aexit__(None, None, None) - raise - sb = cand - break - except Exception as e: - last_err = e - logger.warning( - "[coding_agent_rl] provision attempt %d/%d failed: %s: %s", - attempt + 1, - SWE_BOOT_RETRIES, - type(e).__name__, - str(e)[:200], - ) - await asyncio.sleep(1 + attempt) - if sb is None: - assert last_err is not None - raise last_err - try: - yield sb - finally: - await sb.__aexit__(None, None, None) - - -async def install_node22(sb: Sandbox, host_tarball: Path) -> None: - """Node 22 over the base image (Debian 12 ships 16; cli.js needs >= 20). - Decompresses .xz on the host (cached) so sandboxes without xz-utils can - still run plain `tar xf`. npm prefix=/usr/local required for sweap-images.""" - host_tarball = Path(host_tarball) - if host_tarball.suffix == ".xz": - plain = Path(tempfile.gettempdir()) / f"coding_agent_rl.{host_tarball.stem}.tar" - if not plain.exists(): - tmp = plain.with_suffix(".tar.partial") - with lzma.open(host_tarball, "rb") as src, open(tmp, "wb") as dst: - shutil.copyfileobj(src, dst) - os.replace(tmp, plain) - host_tarball = plain - await sb.write_file("/tmp/node22.tar", host_tarball) - await sb.exec( - "set -e && mkdir -p /opt/node22 && " - "tar xf /tmp/node22.tar -C /opt/node22 --strip-components=1 && " - "ln -sf /opt/node22/bin/node /usr/local/bin/node && " - "ln -sf /opt/node22/bin/npm /usr/local/bin/npm && " - "ln -sf /opt/node22/bin/npx /usr/local/bin/npx && " - "hash -r 2>/dev/null || true && node --version && npm --version", - user="root", - timeout=180, - check=True, - ) - - -async def install_claude_code(sb: Sandbox, host_tarball: Path) -> None: - await sb.write_file("/tmp/claude-code.tgz", host_tarball) - await sb.exec( - "npm install -g --prefix=/usr/local --no-audit --no-fund /tmp/claude-code.tgz " - "&& ls -la /usr/local/bin/claude && /usr/local/bin/claude --version", - user="root", - timeout=300, - check=True, - ) - - -async def ensure_agent_user(sb: Sandbox, workdir: str) -> None: - """Create the unprivileged 'agent' user that owns workdir + can git diff. - Settings file pre-acks bypass-permissions so claude-code starts headless.""" - await sb.exec( - f"id agent >/dev/null 2>&1 || useradd -m -s /bin/bash agent && " - f"chown -R agent:agent /home/agent {workdir} && " - f"git config --system --add safe.directory '*' && id agent && " - f"mkdir -p /home/agent/.claude && " - f'echo \'{{"hasCompletedOnboarding": true, "bypassPermissionsModeAccepted": true}}\' ' - f"| tee /home/agent/.claude.json /home/agent/.claude/settings.json > /dev/null && " - f"chown -R agent:agent /home/agent/.claude /home/agent/.claude.json", - user="root", - check=True, - timeout=60, - ) - - -async def apply_before_repo_set_cmd(sb: Sandbox, workdir: str, swepro: dict) -> None: - """Run swepro['before_repo_set_cmd'] in the sandbox if present (no-op if not).""" - before = swepro.get("before_repo_set_cmd") if swepro else None - if not before: - return - payload = f"set -e\ncd {workdir}\n{before}\n" - await sb.exec( - "mkdir -p /workspace/swepro_setup && chown agent:agent /workspace/swepro_setup", user="root", check=True - ) - await sb.write_file("/workspace/swepro_setup/before.sh", payload, user="agent") - await sb.exec("bash /workspace/swepro_setup/before.sh", user="agent", check=False, timeout=600) - - -# --------------------------------------------------------------------------- -# Agent run (workspace prep + claude-code spawn + done-marker poll) -# --------------------------------------------------------------------------- -async def run_claude_code( - sb: Sandbox, - *, - workdir: str, - session_id: str, - adapter_url: str, - time_budget_sec: int, - problem_statement: str = "", - swepro: dict | None = None, - pre_commands: list[str] | str | None = None, - prompt: str | None = None, -) -> int: - """Prepare the SWE workspace, write PROBLEM_STATEMENT.md, then run CC.""" - await ensure_agent_user(sb, workdir) - if swepro: - await apply_before_repo_set_cmd(sb, workdir, swepro) - if pre_commands: - await apply_pre_commands(sb, workdir, pre_commands) - await sb.write_file( - f"{workdir}/PROBLEM_STATEMENT.md", - problem_statement or "", - user="agent", - ) - return await _spawn_claude_code( - sb, - workdir=workdir, - session_id=session_id, - adapter_url=adapter_url, - prompt=prompt or CC_PROMPT, - time_budget_sec=time_budget_sec, - ) - - -async def _spawn_claude_code( - sb: Sandbox, - *, - workdir: str, - session_id: str, - adapter_url: str, - prompt: str, - time_budget_sec: int, -) -> int: - """Spawn claude-code detached + poll a done-marker file. - - E2B's gateway resets HTTP/2 around 6.5 min, so we can't keep a long-lived - foreground exec. The launcher writes the exit code into a marker file - and we poll it every 5s via short RPCs (which also keeps the sandbox - alive against idle GC).""" - done = f"{workdir}/.cagent_done" - launcher = f"{workdir}/.cagent_run.sh" - traj = f"{workdir}/claude_code_trajectory.jsonl" - - launcher_body = ( - "#!/bin/bash\n" - f"cd {workdir}\n" - "export HOME=/home/agent\n" - f"/usr/local/bin/claude -p {json.dumps(prompt)} " - f"--permission-mode bypassPermissions " - f"--output-format stream-json --include-partial-messages " - f"--include-hook-events --verbose " - f"{os.environ.get('SWE_CLAUDE_EXTRA_ARGS', '').strip()} " - f"2>&1 | tee {shlex.quote(traj)}\n" - f"echo $? > {done}\n" - ) - await sb.write_file(launcher, launcher_body, user="agent") - await sb.exec(f"chmod +x {launcher}", user="agent", timeout=30) - - env = { - "ANTHROPIC_BASE_URL": adapter_url, - "ANTHROPIC_AUTH_TOKEN": session_id, - "ANTHROPIC_MODEL": "vime-actor", - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", - "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", - } - env_keys = ",".join(env.keys()) - await sb.exec( - f"runuser -u agent --whitelist-environment={env_keys}" - f" -- bash -c 'setsid {launcher} < /dev/null > /dev/null 2>&1 &'", - user="root", - env=env, - timeout=30, - check=True, - ) - - deadline = time.time() + time_budget_sec - exit_code = -2 # convention: -2 = budget exceeded - while time.time() < deadline: - await asyncio.sleep(5) - ec, out, _ = await sb.exec( - f"test -f {done} && cat {done}", - user="agent", - timeout=15, - check=False, - ) - if ec == 0: - try: - exit_code = int((out or "").strip() or "-1") - except ValueError: - exit_code = -1 - break - return exit_code - - -async def git_diff(sb: Sandbox, workdir: str) -> str: - cmd = ( - f"cd {workdir} && git add -N . && " - f"git diff -- . ':(exclude)PROBLEM_STATEMENT.md' " - f"':(exclude)claude_code_trajectory.jsonl' " - f"':(exclude).cagent_done' ':(exclude).cagent_run.sh'" - ) - _, out, _ = await sb.exec(cmd, user="agent", timeout=120) - return out - - -# --------------------------------------------------------------------------- -# Eval (fresh sandbox, apply diff, run dataset tests) -# --------------------------------------------------------------------------- -async def evaluate( - *, - image: str, - workdir: str, - diff_text: str, - swepro: dict | None = None, - eval_cmd: str | None = None, - pre_commands: list[str] | str | None = None, - timeout_sec: int = 600, -) -> tuple[float, bool, bool]: - """Returns (reward, solved, applied_cleanly). - - No-test-cheating guarantee: the eval sandbox is built from the same image - but starts CLEAN, so only the model-produced diff affects reward.""" - if not (swepro or eval_cmd): - logger.warning("[e2b.evaluate] no swepro/eval_cmd; reward=0") - return 0.0, False, True - - async with E2BSandbox(image) as ev: - await ensure_agent_user(ev, workdir) - if swepro: - await _setup_swepro_assets(ev, swepro) - await apply_before_repo_set_cmd(ev, workdir, swepro) - if pre_commands: - await apply_pre_commands(ev, workdir, pre_commands) - - applied = await _apply_diff(ev, workdir, diff_text) - if not applied: - return 0.0, False, False - - if swepro: - r, s = await _run_swepro(ev, workdir, swepro, timeout_sec) - return r, s, True - r, s = await _run_eval_cmd(ev, workdir, eval_cmd, timeout_sec) - return r, s, True - - -async def _setup_swepro_assets(ev: Sandbox, swepro: dict) -> None: - await ev.exec(f"mkdir -p {_SWEPRO_DIR} && chmod 777 {_SWEPRO_DIR}", user="root", check=True) - for k, dst in [("run_script_path", "run_script.sh"), ("parser_script_path", "parser.py")]: - host_p = swepro.get(k) - if host_p: - text = Path(host_p).read_text() - await ev.write_file(f"{_SWEPRO_DIR}/{dst}", text, user="root") - await ev.exec(f"chmod 755 {_SWEPRO_DIR}/* && chown -R agent:agent {_SWEPRO_DIR}", user="root", check=True) - - -async def apply_pre_commands(ev: Sandbox, workdir: str, pre: list[str] | str) -> None: - # Public: also called by generate.py to keep the work sandbox baseline - # aligned with eval (sweb-style pre_commands typically `git checkout - # -f`, so skipping in work sandbox makes the model's diff - # context mismatch the eval base -> 100% apply failure). - if isinstance(pre, str): - body = pre.replace("\\n", "\n") - else: - body = "\n".join(c for c in (pre or []) if c) - await ev.write_file(_PRE, "set -e\n" + body, user="agent") - await ev.exec(f"chmod 755 {_PRE} && cd {workdir} && bash {_PRE}", user="agent", check=False, timeout=600) - - -async def _apply_diff(ev: Sandbox, workdir: str, diff_text: str) -> bool: - if not diff_text.strip(): - return True - await ev.write_file(_PATCH, diff_text, user="agent") - for cmd in [ - f"cd {workdir} && git apply --3way --whitespace=nowarn {_PATCH}", - f"cd {workdir} && git apply --whitespace=nowarn {_PATCH}", - f"cd {workdir} && patch -p1 --no-backup-if-mismatch < {_PATCH}", - ]: - ec, _, _ = await ev.exec(cmd, user="agent", check=False, timeout=120) - if ec == 0: - return True - return False - - -async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict, timeout: int) -> tuple[float, bool]: - test_arg = ",".join(swepro.get("selected_test_files") or []) - stdout_f = f"{_SWEPRO_DIR}/stdout.log" - stderr_f = f"{_SWEPRO_DIR}/stderr.log" - result_f = f"{_SWEPRO_DIR}/result.json" - await ev.exec( - f"cd {workdir} && bash {_SWEPRO_DIR}/run_script.sh " - f"{json.dumps(test_arg)} > {stdout_f} 2> {stderr_f} || true", - user="agent", - check=False, - timeout=timeout, - ) - await ev.exec( - f"python3 {_SWEPRO_DIR}/parser.py {stdout_f} {stderr_f} {result_f}", - user="agent", - check=False, - timeout=120, - ) - raw = await ev.read_file(result_f, user="agent") - parsed = json.loads(raw) if raw else {"tests": []} - passed = {t["name"] for t in parsed.get("tests", []) if t.get("status") == "PASSED"} - required = set(swepro.get("fail_to_pass") or []) | set(swepro.get("pass_to_pass") or []) - solved = bool(required) and required.issubset(passed) - return (1.0 if solved else 0.0), solved - - -async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> tuple[float, bool]: - ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) - return (1.0 if ec == 0 else 0.0), ec == 0 diff --git a/examples/coding_agent_rl/swe.py b/examples/coding_agent_rl/swe.py new file mode 100644 index 000000000..0ada75186 --- /dev/null +++ b/examples/coding_agent_rl/swe.py @@ -0,0 +1,492 @@ +"""SWE task layer: dataset parsing, workspace prep, diff capture, fresh-sandbox eval. + +One module, two grading protocols selected per-call (never an import-time side +effect): + + - "scaleswe" (default): scaleswe data shape (image_url + pre_commands + + swepro/eval_cmd/f2p_script); custom "exit 0 == solved" grading. + - "swebench": SWE-bench Verified (remote_env_info.{image,base_commit, + test_patch,FAIL_TO_PASS,PASS_TO_PASS,version}); graded with swebench's + official make_test_spec + get_eval_report so each repo uses its own + test_cmd and log parser. + +The only thing that varies by protocol is the dataset schema and how a +diff is scored. Everything sandbox-side (prepare_workspace / git_diff / +apply_diff / pre_commands) is shared and lives here once. +``get_metadata(sample, protocol)`` produces the ``md`` dict; the +protocol-specific grading payload is carried under ``md["grading"]`` +and is opaque to generate.py (which only reads instance_id / image / workdir). + +Harness-agnostic on purpose -- nothing here is Claude-specific. ``SWE_PROMPT`` is +the task instruction (semantics, not CLI syntax). The only place a task meets a +harness is the prompt, which the orchestrator passes into ``harness.run()``. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import tempfile +from pathlib import Path +from typing import Any, NamedTuple + +from vime.agent import sandbox as agent_sandbox +from vime.agent.adapters.common import flatten_content +from vime.agent.sandbox import E2BSandbox, Sandbox, exec_and_wait +from vime.utils.types import Sample + +try: + from swebench.harness.grading import get_eval_report # type: ignore + from swebench.harness.test_spec.test_spec import make_test_spec # type: ignore + + _SWEBENCH_IMPORT_ERROR: Exception | None = None +except Exception as _exc: # pragma: no cover - import-time diagnostic + get_eval_report = None # type: ignore + make_test_spec = None # type: ignore + _SWEBENCH_IMPORT_ERROR = _exc + +logger = logging.getLogger(__name__) + +PROTOCOL_SCALESWE = "scaleswe" +PROTOCOL_SWEBENCH = "swebench" + +# Paths inside the sandbox (avoid clashes with image-shipped paths). +_PATCH = "/workspace/__cagent_patch__.diff" +_PRE = "/workspace/__cagent_pre__.sh" +_F2P = "/workspace/__cagent_f2p__.py" +_SWEPRO_DIR = "/workspace/swepro_eval" + +SWE_PROMPT = os.environ.get( + "SWE_CC_PROMPT", + "Read PROBLEM_STATEMENT.md in the current directory and resolve the issue. " + "Edit source files only (do NOT touch tests). After editing, run the relevant " + "tests to verify your fix passes. Do NOT modify PROBLEM_STATEMENT.md and do " + "NOT commit. When finished, print a one-line summary and exit.", +) + + +class EvalResult(NamedTuple): + """Grading outcome. Tuple-compatible: ``reward, applied = run_evaluation(...)``.""" + + reward: float + applied_cleanly: bool + + +def get_metadata(sample: Sample, protocol: str = PROTOCOL_SCALESWE) -> dict[str, Any]: + if protocol == PROTOCOL_SWEBENCH: + return _metadata_swebench(sample) + return _metadata_scaleswe(sample) + + +def _metadata_scaleswe(sample: Sample) -> dict[str, Any]: + """scaleswe shape: flat ``metadata.*`` (+ a few ``remote_env_info`` fallbacks). + + ``f2p_script`` (a self-contained pytest file ending in + ``sys.exit(pytest.main(...))``) is carried verbatim; the grader materializes + and runs it via ``write_file`` so no shell-quoting workaround is needed here. + """ + m = sample.metadata or {} + rem = m.get("remote_env_info") or {} + label = sample.label if (isinstance(sample.label, str) and len(sample.label) < 256) else None + swepro = m.get("swepro") + eval_cmd = m.get("eval_cmd") + f2p_script = rem.get("f2p_script") + looks_swebench = bool(rem.get("test_patch")) and not (swepro or eval_cmd or f2p_script) + return { + "protocol": PROTOCOL_SCALESWE, + "instance_id": m.get("instance_id") or rem.get("instance_id") or label or "unknown", + "image": m.get("image") or rem.get("image_url"), + "workdir": m.get("workdir") or rem.get("workdir"), + "problem_statement": m.get("problem_statement") or _coerce_prompt(sample.prompt), + "looks_swebench": looks_swebench, + "grading": { + "swepro": swepro, + "eval_cmd": eval_cmd, + "f2p_script": f2p_script, + "pre_commands": m.get("pre_commands") or rem.get("pre_commands"), + }, + } + + +def _metadata_swebench(sample: Sample) -> dict[str, Any]: + """SWE-bench Verified shape: carry the full instance dict through so + make_test_spec gets every field it needs (version, hints_text, ...).""" + m = sample.metadata or {} + rem = m.get("remote_env_info") or {} + instance = { + "instance_id": rem.get("instance_id") or "unknown", + "repo": rem.get("repo") or "", + "version": rem.get("version"), + "base_commit": rem.get("base_commit") or "", + "problem_statement": rem.get("problem_statement") or _coerce_prompt(sample.prompt), + "hints_text": rem.get("hints_text") or "", + "test_patch": rem.get("test_patch") or "", + "FAIL_TO_PASS": rem.get("FAIL_TO_PASS"), + "PASS_TO_PASS": rem.get("PASS_TO_PASS"), + "environment_setup_commit": rem.get("environment_setup_commit"), + } + return { + "protocol": PROTOCOL_SWEBENCH, + "instance_id": instance["instance_id"], + "image": rem.get("image"), + "workdir": rem.get("workdir") or "/testbed", + "problem_statement": instance["problem_statement"], + "grading": {"sweb_instance": instance}, + } + + +def _coerce_prompt(prompt) -> str: + """Extract the user-message text from a prompt (str or chat-message list).""" + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list): + for m in prompt: + if isinstance(m, dict) and m.get("role") == "user": + return flatten_content(m.get("content")) + return "" + + +def evaluability_check(md: dict) -> str | None: + if md.get("protocol") == PROTOCOL_SWEBENCH: + return _evaluability_check_swebench(md) + return "protocol_row_mismatch:looks_swebench" if md.get("looks_swebench") else None + + +def _evaluability_check_swebench(md: dict) -> str | None: + if _SWEBENCH_IMPORT_ERROR is not None: + return f"swebench_import_failed:{type(_SWEBENCH_IMPORT_ERROR).__name__}" + inst = md.get("grading", {}).get("sweb_instance") or {} + if not inst.get("repo"): + return "missing_repo" + if not inst.get("base_commit"): + return "missing_base_commit" + if not (inst.get("test_patch") or "").strip(): + return "missing_test_patch" + try: + _ = _build_test_spec(inst).eval_script # surfaces per-repo construction errors here, not later + except Exception as e: # KeyError on unknown repo/version, etc. + return f"make_test_spec_failed:{type(e).__name__}" + return None + + +# --------------------------------------------------------------------------- +# Workspace prep (agent sandbox, before harness.run) +# --------------------------------------------------------------------------- +async def prepare_workspace(sb: Sandbox, workdir: str, md: dict) -> None: + """Prep the agent sandbox, then drop PROBLEM_STATEMENT.md. + + Assumes the agent user already owns ``workdir`` (the harness's ``run()`` calls + ``ensure_agent_user``; the orchestrator runs this before ``run()`` and the + agent user is created lazily there). To stay independent of call order we + create the agent user here too -- it is idempotent. + """ + await agent_sandbox.ensure_agent_user(sb, workdir) + if md.get("protocol") == PROTOCOL_SCALESWE: + grading = md.get("grading") or {} + swepro = grading.get("swepro") + if swepro: + await apply_before_repo_set_cmd(sb, workdir, swepro) + pre_commands = grading.get("pre_commands") + if pre_commands: + await apply_pre_commands(sb, workdir, pre_commands) + await sb.write_file( + f"{workdir}/PROBLEM_STATEMENT.md", + md.get("problem_statement") or "", + user="agent", + ) + + +async def apply_before_repo_set_cmd(sb: Sandbox, workdir: str, swepro: dict) -> None: + """Run swepro['before_repo_set_cmd'] in the sandbox if present (no-op if not).""" + before = swepro.get("before_repo_set_cmd") + if not before: + return + payload = f"set -e\ncd {workdir}\n{before}\n" + await sb.exec( + "mkdir -p /workspace/swepro_setup && chown agent:agent /workspace/swepro_setup", user="root", check=True + ) + await sb.write_file("/workspace/swepro_setup/before.sh", payload, user="agent") + await sb.exec("bash /workspace/swepro_setup/before.sh", user="agent", check=False, timeout=600) + + +async def apply_pre_commands(sb: Sandbox, workdir: str, pre: list[str] | str) -> None: + # Public: also called for the work sandbox to keep its baseline aligned with + # eval (sweb-style pre_commands typically `git checkout -f`, so + # skipping in the work sandbox makes the model's diff context mismatch the + # eval base -> 100% apply failure). + if isinstance(pre, str): + body = pre.replace("\\n", "\n") + else: + body = "\n".join(c for c in (pre or []) if c) + await sb.write_file(_PRE, "set -e\n" + body, user="agent") + await sb.exec(f"chmod 755 {_PRE} && cd {workdir} && bash {_PRE}", user="agent", check=False, timeout=600) + + +# --------------------------------------------------------------------------- +# Diff capture (agent sandbox, after harness.run) +# --------------------------------------------------------------------------- +async def git_diff(sb: Sandbox, workdir: str) -> str: + cmd = f"cd {workdir} && git add -N . && git diff -- . ':(exclude)PROBLEM_STATEMENT.md' ':(exclude).harness/'" + _, out, _ = await sb.exec(cmd, user="agent", timeout=120) + return out + + +# --------------------------------------------------------------------------- +# Eval dispatch (fresh sandbox, apply diff, run dataset tests) +# --------------------------------------------------------------------------- +async def run_evaluation(md: dict, *, diff_text: str, timeout_sec: int) -> EvalResult: + """Uniform entry point: dispatch to the protocol's grader. + + No-test-cheating guarantee (both grading protocols): the eval sandbox is built from + the same image but starts CLEAN, so only the model-produced diff affects + reward.""" + if md.get("protocol") == PROTOCOL_SWEBENCH: + return await _grade_swebench(md, diff_text, timeout_sec) + return await _grade_scaleswe(md, diff_text, timeout_sec) + + +# --------------------------------------------------------------------------- +# scaleswe grader +# --------------------------------------------------------------------------- +async def _grade_scaleswe(md: dict, diff_text: str, timeout_sec: int) -> EvalResult: + """Three mutually-exclusive grading paths, in priority order: swepro test + harness, a shell ``eval_cmd``, or a self-contained ``f2p_script`` pytest + file. All resolve to "exit 0 == solved", reward is 1.0 iff solved.""" + image = md["image"] + workdir = md["workdir"] + grading = md.get("grading") or {} + swepro = grading.get("swepro") + eval_cmd = grading.get("eval_cmd") + f2p_script = grading.get("f2p_script") + pre_commands = grading.get("pre_commands") + + if not (swepro or eval_cmd or f2p_script): + logger.warning("[swe.scaleswe] no swepro/eval_cmd/f2p_script; reward=0") + return EvalResult(0.0, True) + + async with E2BSandbox(image) as ev: + await agent_sandbox.ensure_agent_user(ev, workdir) + if swepro: + await _setup_swepro_assets(ev, swepro) + await apply_before_repo_set_cmd(ev, workdir, swepro) + if pre_commands: + await apply_pre_commands(ev, workdir, pre_commands) + + applied = await _apply_diff(ev, workdir, diff_text) + if not applied: + return EvalResult(0.0, False) + + if swepro: + r = await _run_swepro(ev, workdir, swepro, timeout_sec) + elif eval_cmd: + r = await _run_eval_cmd(ev, workdir, eval_cmd, timeout_sec) + else: + r = await _run_f2p_script(ev, workdir, f2p_script, timeout_sec) + return EvalResult(r, True) + + +async def _setup_swepro_assets(ev: Sandbox, swepro: dict) -> None: + await ev.exec(f"mkdir -p {_SWEPRO_DIR} && chmod 777 {_SWEPRO_DIR}", user="root", check=True) + for k, dst in [("run_script_path", "run_script.sh"), ("parser_script_path", "parser.py")]: + host_p = swepro.get(k) + if host_p: + await ev.write_file(f"{_SWEPRO_DIR}/{dst}", Path(host_p), user="root") + await ev.exec(f"chmod 755 {_SWEPRO_DIR}/* && chown -R agent:agent {_SWEPRO_DIR}", user="root", check=True) + + +async def _apply_diff(ev: Sandbox, workdir: str, diff_text: str) -> bool: + if not diff_text.strip(): + return True + await ev.write_file(_PATCH, diff_text, user="agent") + # First-success-wins ladder collapsed into one exec (one sandbox round-trip). + ladder = " || ".join( + f"({cmd})" + for cmd in ( + f"git apply --3way --whitespace=nowarn {_PATCH}", + f"git apply --whitespace=nowarn {_PATCH}", + f"patch -p1 --no-backup-if-mismatch < {_PATCH}", + ) + ) + ec, _, _ = await ev.exec(f"cd {workdir} && ({ladder})", user="agent", check=False, timeout=120) + return ec == 0 + + +async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict, timeout: int) -> float: + test_arg = ",".join(swepro.get("selected_test_files") or []) + stdout_f = f"{_SWEPRO_DIR}/stdout.log" + stderr_f = f"{_SWEPRO_DIR}/stderr.log" + result_f = f"{_SWEPRO_DIR}/result.json" + await ev.exec( + f"cd {workdir} && bash {_SWEPRO_DIR}/run_script.sh {json.dumps(test_arg)} > {stdout_f} 2> {stderr_f} || true", + user="agent", + check=False, + timeout=timeout, + ) + await ev.exec( + f"python3 {_SWEPRO_DIR}/parser.py {stdout_f} {stderr_f} {result_f}", + user="agent", + check=False, + timeout=120, + ) + raw = await ev.read_file(result_f, user="agent") + parsed = json.loads(raw) if raw else {"tests": []} + passed = {t["name"] for t in parsed.get("tests", []) if t.get("status") == "PASSED"} + required = set(swepro.get("fail_to_pass") or []) | set(swepro.get("pass_to_pass") or []) + solved = bool(required) and required.issubset(passed) + return 1.0 if solved else 0.0 + + +async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> float: + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) + return 1.0 if ec == 0 else 0.0 + + +async def _run_f2p_script(ev: Sandbox, workdir: str, script: str, timeout: int) -> float: + # sweb f2p_script is a self-contained pytest file ending in + # `sys.exit(pytest.main([...]))`; write it verbatim (no shell quoting) and + # let python's exit code carry the pass/fail signal. + await ev.write_file(_F2P, script, user="agent") + ec, _, _ = await ev.exec(f"cd {workdir} && python {_F2P}", user="agent", check=False, timeout=timeout) + return 1.0 if ec == 0 else 0.0 + + +# Mirror of swebench.harness.run_evaluation.GIT_APPLY_CMDS: try each in order, +# first success wins. The `patch --fuzz` tier rescues diffs `git apply` rejects. +_GIT_APPLY_CMDS = ( + "git apply --verbose", + "git apply --verbose --reject", + "patch --batch --fuzz=5 -p1 -i", +) + + +async def _apply_model_patch(ev: Sandbox, workdir: str) -> bool: + """Apply /tmp/patch.diff via the GIT_APPLY_CMDS ladder; True if applied + (or empty). Empty patch is a no-op success -- eval then scores it 0 on its + own (no source change -> tests still fail).""" + ladder = " || ".join(f"{cmd} /tmp/patch.diff" for cmd in _GIT_APPLY_CMDS) + cmd = ( + f"cd {workdir} && git config --global --add safe.directory {workdir} " + f"&& if [ -s /tmp/patch.diff ]; then {ladder}; fi" + ) + ec, _, _ = await ev.exec(cmd, user="root", check=False, timeout=120) + return ec == 0 + + +def _build_test_spec(inst: dict): + """make_test_spec(inst). Shared by evaluability_check and the grader; may + raise (KeyError on unknown repo/version).""" + return make_test_spec(inst) # type: ignore[misc] + + +def _eval_report_from_log(ts, instance_id: str, diff_text: str, log: str) -> dict: + """Run swebench's get_eval_report against the captured test log. It reads + from a file path, so write the log to a tempfile, parse, and clean up.""" + tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".log", delete=False) + try: + tmp.write(log) + tmp.flush() + tmp.close() + prediction = { + "instance_id": instance_id, + "model_patch": diff_text or "", + "model_name_or_path": "swe", + } + return get_eval_report( # type: ignore[misc] + test_spec=ts, + prediction=prediction, + test_log_path=tmp.name, + include_tests_status=True, + ) + finally: + try: + os.unlink(tmp.name) + except OSError: + pass + + +def _ratio(d: dict) -> tuple[int, int]: + """(passed, total) from a {success: [...], failure: [...]} bucket.""" + passed, failed = d.get("success", []), d.get("failure", []) + return len(passed), len(passed) + len(failed) + + +def _log_swebench_result(instance_id: str, exit_code, info: dict, log: str) -> None: + """Emit the per-instance grading outcome with test-bucket ratios; on a + non-resolved row that parsed NO test lines, surface the log tail so failures + (missing pytest plugin, conda not activated, ...) can be diagnosed.""" + if info.get("resolved"): + logger.info("[swe.swebench] %s: reward=1 exit_code=%s", instance_id, exit_code) + return + ts_status = info.get("tests_status") or {} + f2p_pass, f2p_total = _ratio(ts_status.get("FAIL_TO_PASS", {})) + p2p_pass, p2p_total = _ratio(ts_status.get("PASS_TO_PASS", {})) + nothing_parsed = not (f2p_total or p2p_total) + tail = log[-800:] if nothing_parsed else "" + logger.info( + "[swe.swebench] %s: reward=0 exit_code=%s patch_applied=%s F2P=(%d/%d) P2P=(%d/%d)%s", + instance_id, + exit_code, + bool(info.get("patch_successfully_applied")), + f2p_pass, + f2p_total, + p2p_pass, + p2p_total, + f" tail={tail!r}" if tail else "", + ) + + +async def _grade_swebench(md: dict, diff_text: str, timeout_sec: int) -> EvalResult: + """reward=1.0 iff sweb's get_eval_report declares the instance ``resolved``.""" + instance_id = md["instance_id"] + inst = md["grading"]["sweb_instance"] + + if _SWEBENCH_IMPORT_ERROR is not None: + logger.error( + "[swe.swebench] %s: swebench import failed: %r; reward=0", + instance_id, + _SWEBENCH_IMPORT_ERROR, + ) + return EvalResult(0.0, True) + + try: + ts = _build_test_spec(inst) + eval_sh = ts.eval_script # may raise on unknown repo/version + except Exception as e: + logger.warning("[swe.swebench] %s: make_test_spec/eval_script failed: %s; reward=0", instance_id, e) + return EvalResult(0.0, True) + + image = md["image"] + if not image: + logger.warning("[swe.swebench] %s: missing image; reward=0", instance_id) + return EvalResult(0.0, True) + + async with E2BSandbox(image) as ev: + await asyncio.gather( + ev.write_file("/tmp/patch.diff", diff_text or "", user="root"), + ev.write_file("/tmp/eval.sh", eval_sh, user="root"), + ) + # Apply the model patch first (eval_script assumes it is already applied); + # if no apply strategy works, the instance is unsolvable -- skip the eval. + if not await _apply_model_patch(ev, md["workdir"]): + logger.warning("[swe.swebench] %s: model patch failed to apply; reward=0", instance_id) + return EvalResult(0.0, False) + exit_code, log = await exec_and_wait( + ev, cmd="bash /tmp/eval.sh", user="root", time_budget_sec=timeout_sec, tag="eval", want_output=True + ) + + try: + report = _eval_report_from_log(ts, instance_id, diff_text, log) + except Exception as e: + logger.warning( + "[swe.swebench] %s: get_eval_report failed: %s; reward=0 (tail=%r)", + instance_id, + e, + log[-600:], + ) + return EvalResult(0.0, True) + + info = report.get(instance_id, {}) + _log_swebench_result(instance_id, exit_code, info, log) + return EvalResult(1.0 if info.get("resolved") else 0.0, bool(info.get("patch_successfully_applied"))) diff --git a/examples/delta_weight_sync/README.md b/examples/delta_weight_sync/README.md new file mode 100644 index 000000000..477dd84e2 --- /dev/null +++ b/examples/delta_weight_sync/README.md @@ -0,0 +1,41 @@ +# Delta Weight Sync + +Non-colocated weight sync that ships only the **changed bytes** between two syncs instead of a +full checkpoint, for training/inference disaggregation across clusters or datacenters. The +trainer publishes per-tensor deltas to a shared filesystem as a canonical HF checkpoint +directory; each engine's `/pull_weights` applies them into a host-local checkpoint on every +host it spans, and the engines reload through the ordinary `update_weights_from_disk` path — +vime only ever talks to one endpoint per engine. + +See [Delta Weight Sync](../../docs/en/advanced/delta-weight-sync.md) for the full mechanism, +encodings, integrity checks, and shared-filesystem visibility hooks. + +## Try it + +`run-glm4.7-30B-A3B-delta.sh` runs the disk delta path on GLM-4.7-Flash, non-colocated across a +2-node (16-GPU) Ray cluster. See its header for prerequisites. + +## Minimal flags + +Add to a non-colocated training run (the trainer and engines only need to share the filesystem +at `--update-weight-disk-dir`): + +```bash +--update-weight-mode delta \ +--update-weight-transport disk \ +--update-weight-disk-dir /shared/fs/delta-updates \ +--update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt \ +--update-weight-delta-encoding xor \ +--update-weight-delta-checksum xxh3-128 +``` + +- `--update-weight-disk-dir` — shared directory the trainer writes deltas to and the hosts read. +- `--update-weight-local-checkpoint-dir` — host-local full HF checkpoint the delta patches in + place; materialized from the engine's model path on the first `/pull_weights`. +- `--update-weight-delta-encoding` — `xor` (smallest/fastest) or `overwrite` (idempotent). +- `--update-weight-delta-checksum` — `xxh3-128` (default), `blake3`, or `adler32`. + +For object-store-backed volumes that need an explicit commit/refresh to make writes visible +across hosts, supply `--custom-update-weight-post-write-path` (trainer side) / +`--vllm-custom-pull-weights-pre-read-hook` (engine side) — no vendor-specific code lives in vime +or vllm; see the doc. diff --git a/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh b/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh new file mode 100644 index 000000000..ebac2a9eb --- /dev/null +++ b/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh @@ -0,0 +1,109 @@ +#!/bin/bash +# Disk delta weight-sync demo on GLM-4.7-Flash (30B-A3B), non-colocated, 2 nodes x 8 GPU. +# The trainer publishes per-tensor deltas to --update-weight-disk-dir as a canonical HF directory; +# each engine's /pull_weights applies them into --update-weight-local-checkpoint-dir on every host +# it spans, and the engine reloads via the vanilla update_weights_from_disk path. +# +# Prerequisites: +# - A 2-node (16-GPU) Ray cluster, this script run on the head node. +# - GLM-4.7-Flash HF checkpoint + its torch_dist conversion (tools/convert_hf_to_torch_dist.py). +# - dapo-math-17k.jsonl. +# - --update-weight-disk-dir on a filesystem both nodes share. On an object-store-backed volume +# that needs an explicit commit/refresh to surface writes across hosts, also pass +# --custom-update-weight-post-write-path / --vllm-custom-pull-weights-pre-read-hook (see the doc). + +set -ex +export PYTHONUNBUFFERED=1 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/glm4.7-30B-A3B.sh" + +MODEL_DIR=${MODEL_DIR:-/root/models/GLM-4.7-Flash} +DATA_PATH=${DATA_PATH:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl} + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_DIR}" + --ref-load "${MODEL_DIR}_torch_dist" +) + +ROLLOUT_ARGS=( + --prompt-data "${DATA_PATH}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3 + --rollout-batch-size 32 + --n-samples-per-prompt 4 + --rollout-max-response-len 8192 + --global-batch-size 128 +) + +# Disk delta weight sync (the point of this example). +WEIGHT_SYNC_ARGS=( + --update-weight-mode delta + --update-weight-transport disk + --update-weight-disk-dir /shared/fs/glm47-delta-updates + --update-weight-local-checkpoint-dir /local/nvme/glm47-rollout-ckpt + --update-weight-delta-encoding xor + --update-weight-delta-checksum xxh3-128 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --pipeline-model-parallel-size 2 + --context-parallel-size 2 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + --sequence-parallel + --use-dynamic-batch-size + --max-tokens-per-gpu 32768 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.0 + --kl-loss-type low_var_kl +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.8 + --vllm-data-parallel-size 8 + --vllm-enable-expert-parallel +) + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" + } +}" + +# Non-colocated: 16 actor GPUs (2 x 8) train while a 16-GPU rollout pool generates (delta mode +# requires non-colocation). +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 2 \ + --actor-num-gpus-per-node 8 \ + --rollout-num-gpus 16 \ + ${MODEL_ARGS[@]} \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${WEIGHT_SYNC_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${VLLM_ARGS[@]}" diff --git a/examples/dspark/README.md b/examples/dspark/README.md new file mode 100644 index 000000000..f1ff7a1cd --- /dev/null +++ b/examples/dspark/README.md @@ -0,0 +1,139 @@ +# DSpark Speculative Decoding Training + +This example shows how to train a **DSpark** (semi-autoregressive speculative +decoding) draft model alongside the policy model during RL. DSpark drafts +multiple tokens in parallel from intermediate policy hidden states, then vLLM +verifies them in a single forward pass — accelerating rollouts without a +separate draft training pipeline. + +## Key Features + +- **Online draft training**: The draft model is trained jointly with the policy + during RL, so it stays in distribution as the policy updates. +- **Weight sync to vLLM**: Draft weights are synced to vLLM every rollout step + via direct IPC (colocate) or packed tensor transfer (non-colocate), enabling + immediate speculative decoding in the next rollout. +- **Freeze-policy mode**: Optionally freeze the policy and train only the draft + model, useful when the RL signal is weak or policy degradation is a concern. + +## Prerequisites + +1. **Pre-trained DSpark draft checkpoint** (recommended). Training from random + init converges slowly. Pre-train the draft backbone on a supervised corpus + first, then pass the checkpoint via `--dspark-pretrained-model`. + +2. **Policy model** in both HF and torch_dist formats: + +```bash +cd /root/vime +source scripts/models/qwen3-4B.sh + +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-4B \ + --save /root/Qwen3-4B_torch_dist +``` + +3. **Dataset** (e.g., dapo-math-17k): + +```bash +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k +``` + +## Key Arguments + +| Argument | Description | +|----------|-------------| +| `--dspark-pretrained-model` | Path to pre-trained DSpark safetensors. Strongly recommended. | +| `--dspark-block-size` | Number of draft tokens per block (default: 7). | +| `--dspark-num-draft-layers` | Number of decoder layers in the draft backbone (default: 5). | +| `--dspark-target-layer-ids` | Comma-separated policy layer indices to capture hidden states from (default: "1,9,17,25,33"). | +| `--dspark-freeze-policy` | Freeze policy and train only the draft model. | +| `--dspark-ce-loss-alpha` | Weight for cross-entropy loss (default: 0.1). | +| `--dspark-l1-loss-alpha` | Weight for L1/TV loss (default: 0.9). | +| `--dspark-draft-loss-weight` | Weight multiplying draft loss added to policy loss (default: 1.0). | +| `--vllm-speculative-config` | JSON config passed to vLLM. Setting `method: "dspark"` also enables DSpark draft training. | + +## Mode Comparison + +| Mode | Flag | GPU Layout | Weight Sync | +|------|------|------------|-------------| +| **Colocate** | `--colocate` | Train + rollout share the same GPUs | Direct IPC (fastest) | +| **Non-colocate** | (default) | Train and rollout on separate GPU sets | Packed tensor transfer | + +## Running the Example + +### Colocate Mode (Recommended) + +Train and rollout share the same 8 GPUs. The policy model is offloaded to CPU +during rollout, then restored for the next training step. + +```bash +bash examples/dspark/run-qwen3-4B-dspark-colocate.sh +``` + +GPU layout: + +| GPUs | Role | +|------|------| +| 0–7 | Policy Megatron train + vLLM rollout (colocate) | + +### Non-Colocate Mode + +Train and rollout run on separate GPU groups. This avoids the offload overhead +but requires more GPUs. + +```bash +bash examples/dspark/run-qwen3-4B-dspark-non-colocate.sh +``` + +GPU layout: + +| GPUs | Role | +|------|------| +| 0–3 | Policy Megatron train | +| 4–7 | vLLM rollout (with DSpark speculative decoding) | + +## What to Expect + +On Qwen3-4B with 8x A800 GPUs and a pre-trained DSpark draft checkpoint: + +| Metric | Typical Value | +|--------|---------------| +| Draft acceptance rate | 30–40% | +| Mean acceptance length | 3.2–3.8 | +| Rollout speedup | ~2x vs no speculative decoding | +| Weight sync time | ~10s per step | + +## FAQ + +1. **Do I need a pre-trained draft model?** + Strongly recommended. Training from random init requires many more steps to + converge. Pre-train the draft backbone on a supervised corpus, then pass the + checkpoint via `--dspark-pretrained-model`. + +2. **What does `--dspark-freeze-policy` do?** + It freezes the policy model and trains only the draft model. The policy + logits are detached so gradients only flow to the draft. Use this when the + RL signal is weak or when you want to improve speculative decoding without + affecting the policy. + +3. **How are draft weights synced to vLLM?** + Through vLLM's standard draft weight-update session: colocated engines use + IPC and non-colocated engines use NCCL. + +4. **What is `--dspark-block-size`?** + The number of tokens the draft model predicts in parallel per block. Larger + values increase potential speedup but may reduce acceptance rate. The + default (7) works well for most models. + +5. **How do I choose `--dspark-target-layer-ids`?** + These are the policy layer indices from which hidden states are captured as + input to the draft model. For a 36-layer model, `"1,9,17,25,33"` samples + every 8th layer. More target layers = richer draft input but higher cost. + +## References + +1. [DSpark Paper](https://arxiv.org/abs/2505.14269) — Semi-autoregressive speculative decoding. +2. [vLLM Speculative Decoding Docs](https://docs.vllm.ai/en/latest/features/speculative_decoding/) +3. [vime Speculative Decoding Docs](../../docs/en/advanced/speculative-decoding.md) diff --git a/examples/dspark/run-qwen3-4B-dspark-colocate.sh b/examples/dspark/run-qwen3-4B-dspark-colocate.sh new file mode 100644 index 000000000..00fecb4bd --- /dev/null +++ b/examples/dspark/run-qwen3-4B-dspark-colocate.sh @@ -0,0 +1,166 @@ +#!/bin/bash + +# DSpark speculative decoding training — colocate mode (8 GPUs) +# +# Train and rollout share the same 8 GPUs. The policy model is offloaded to CPU +# during rollout, then restored for the next training step. +# +# Prerequisites: +# - Qwen3-4B HF checkpoint + torch_dist conversion +# - Pre-trained DSpark draft checkpoint (model.safetensors) +# - dapo-math-17k dataset +# +# Usage: bash examples/dspark/run-qwen3-4B-dspark-colocate.sh + +set -ex +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +source "/root/vime/scripts/models/qwen3-4B.sh" + +MODEL_DIR=${MODEL_DIR:-/root/Qwen3-4B} +DSPARK_MODEL=${DSPARK_MODEL:-/root/Qwen3-4B-dspark-pretrained} +DATA_PATH=${DATA_PATH:-/root/dapo-math-17k/dapo-math-17k.jsonl} +SAVE_DIR=${SAVE_DIR:-/root/Qwen3-4B_dspark_colocate/} + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_DIR}" + --ref-load "${MODEL_DIR}_torch_dist" + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL:-50}" +) + +ROLLOUT_ARGS=( + --prompt-data "${DATA_PATH}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout "${NUM_ROLLOUT:-200}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-32}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-8}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-4096}" + --rollout-temperature 1 + + --global-batch-size "${GLOBAL_BATCH_SIZE:-256}" + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-8192}" +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.5 + --vllm-speculative-config '{"method":"dspark","model":"'"${DSPARK_MODEL}"'","num_speculative_tokens":7}' +) + +DSPARK_ARGS=( + --dspark-block-size 7 + --dspark-num-draft-layers 5 + --dspark-target-layer-ids 1,9,17,25,33 + --dspark-markov-rank 256 + --dspark-markov-head-type vanilla + --dspark-num-anchors 512 + --dspark-mask-token-id 151669 + --dspark-ce-loss-alpha 0.5 + --dspark-l1-loss-alpha 0.5 + --dspark-confidence-head-alpha 0.1 + --dspark-loss-decay-gamma 4.0 + --dspark-draft-loss-weight 1.0 + --dspark-freeze-policy + --dspark-intermediate-size 9728 + --dspark-pretrained-model "${DSPARK_MODEL}/model.safetensors" +) + +EVAL_ARGS=( + --eval-interval 100 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --no-persist-layer-norm +) + +# Start Ray with all 8 GPUs +ray start --head --port=6379 --num-gpus=8 \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --working-dir /root/vime \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${DSPARK_ARGS[@]} \ + ${MISC_ARGS[@]} + +# Cleanup +ray stop --force +pkill -9 ray 2>/dev/null || true +pkill -9 -f "train.py" 2>/dev/null || true diff --git a/examples/dspark/run-qwen3-4B-dspark-non-colocate.sh b/examples/dspark/run-qwen3-4B-dspark-non-colocate.sh new file mode 100644 index 000000000..e1e0f99f2 --- /dev/null +++ b/examples/dspark/run-qwen3-4B-dspark-non-colocate.sh @@ -0,0 +1,172 @@ +#!/bin/bash + +# DSpark speculative decoding training — non-colocate mode (8 GPUs) +# +# Train and rollout run on separate GPU groups. The policy model stays on GPU +# during rollout (no offload overhead), and draft weights are synced to vLLM +# via packed tensor transfer. +# +# GPU layout: +# GPUs 0-3: Policy Megatron training (TP=4) +# GPUs 4-7: vLLM rollout with DSpark speculative decoding +# +# Prerequisites: +# - Qwen3-4B HF checkpoint + torch_dist conversion +# - Pre-trained DSpark draft checkpoint (model.safetensors) +# - dapo-math-17k dataset +# +# Usage: bash examples/dspark/run-qwen3-4B-dspark-non-colocate.sh + +set -ex +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +source "/root/vime/scripts/models/qwen3-4B.sh" + +MODEL_DIR=${MODEL_DIR:-/root/Qwen3-4B} +DSPARK_MODEL=${DSPARK_MODEL:-/root/Qwen3-4B-dspark-pretrained} +DATA_PATH=${DATA_PATH:-/root/dapo-math-17k/dapo-math-17k.jsonl} +SAVE_DIR=${SAVE_DIR:-/root/Qwen3-4B_dspark_non_colocate/} + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_DIR}" + --ref-load "${MODEL_DIR}_torch_dist" + --save "${SAVE_DIR}" + --save-interval "${SAVE_INTERVAL:-50}" +) + +ROLLOUT_ARGS=( + --prompt-data "${DATA_PATH}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout "${NUM_ROLLOUT:-200}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-32}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-8}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-4096}" + --rollout-temperature 1 + + --global-batch-size "${GLOBAL_BATCH_SIZE:-256}" + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-8192}" +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus 4 + --rollout-num-gpus-per-engine 4 + --vllm-gpu-memory-utilization 0.85 + --vllm-speculative-config '{"method":"dspark","model":"'"${DSPARK_MODEL}"'","num_speculative_tokens":7}' +) + +DSPARK_ARGS=( + --dspark-block-size 7 + --dspark-num-draft-layers 5 + --dspark-target-layer-ids 1,9,17,25,33 + --dspark-markov-rank 256 + --dspark-markov-head-type vanilla + --dspark-num-anchors 512 + --dspark-mask-token-id 151669 + --dspark-ce-loss-alpha 0.5 + --dspark-l1-loss-alpha 0.5 + --dspark-confidence-head-alpha 0.1 + --dspark-loss-decay-gamma 4.0 + --dspark-draft-loss-weight 1.0 + --dspark-freeze-policy + --dspark-intermediate-size 9728 + --dspark-pretrained-model "${DSPARK_MODEL}/model.safetensors" +) + +EVAL_ARGS=( + --eval-interval 100 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --no-persist-layer-norm +) + +# Start Ray with all 8 GPUs +ray start --head --port=6379 --num-gpus=8 \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --working-dir /root/vime \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${DSPARK_ARGS[@]} \ + ${MISC_ARGS[@]} + +# Cleanup +ray stop --force +pkill -9 ray 2>/dev/null || true +pkill -9 -f "train.py" 2>/dev/null || true diff --git a/examples/eval_multi_task/README.md b/examples/eval_multi_task/README.md new file mode 100644 index 000000000..4ac30ebce --- /dev/null +++ b/examples/eval_multi_task/README.md @@ -0,0 +1,12 @@ +# Multi-Task Evaluation Example + +## Configuring `multi_task.yaml` +- `eval.defaults` defines inference parameters shared by every dataset entry. Override them inside an individual dataset block if needed. +- `eval.datasets` enumerates the datasets to evaluate. Each entry should specify: + - `name`: a short identifier that appears in logs and dashboards. + - `path`: the path to the dataset JSONL file. + - `rm_type`: which reward function to use for scoring. + - `n_samples_per_eval_prompt`: how many candidate completions to generate per prompt. + +## IFBench Notes +- When `ifbench` is used, `vime/rollout/rm_hub/ifbench.py` will automatically prepares the scoring environment, so no additional manual setup is required beyond providing the dataset path. diff --git a/examples/eval_multi_task/multi_task.sh b/examples/eval_multi_task/multi_task.sh new file mode 100644 index 000000000..7a0265964 --- /dev/null +++ b/examples/eval_multi_task/multi_task.sh @@ -0,0 +1,149 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" +source "${REPO_ROOT}/scripts/models/qwen3-4B.sh" +EVAL_CONFIG_PATH="${REPO_ROOT}/examples/eval_multi_task/multi_task.yaml" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B + #--hf-checkpoint /root/Qwen3-4B-FP8 + --ref-load /root/Qwen3-4B_torch_dist + --load /root/Qwen3-4B_vime/ + --save /root/Qwen3-4B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-config "${EVAL_CONFIG_PATH}" +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + --use-wandb + --wandb-project eval + --wandb-group multi_task + --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.7 +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 "${REPO_ROOT}/train.py" \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/examples/eval_multi_task/multi_task.yaml b/examples/eval_multi_task/multi_task.yaml new file mode 100644 index 000000000..bad2d6141 --- /dev/null +++ b/examples/eval_multi_task/multi_task.yaml @@ -0,0 +1,17 @@ +eval: + defaults: + max_response_len: 16384 + top_p: 0.7 + datasets: + - name: aime + path: /root/aime-2024/aime-2024.jsonl + rm_type: deepscaler + n_samples_per_eval_prompt: 16 + - name: gpqa # hf download --repo-type dataset zyzshishui0627/gpqa_diamond --local-dir /root/gpqa + path: /root/gpqa/gpqa_eval.jsonl + rm_type: gpqa + n_samples_per_eval_prompt: 2 + - name: ifbench # hf download --repo-type dataset zyzshishui0627/IFBench --local-dir /root/ifbench + path: /root/ifbench/IFBench_eval.jsonl + rm_type: ifbench + n_samples_per_eval_prompt: 1 diff --git a/examples/eval_multi_task/requirements_ifbench.txt b/examples/eval_multi_task/requirements_ifbench.txt new file mode 100644 index 000000000..78f13fac4 --- /dev/null +++ b/examples/eval_multi_task/requirements_ifbench.txt @@ -0,0 +1,6 @@ +emoji +immutabledict +nltk +numpy==1.26.4 +spacy==3.7.4 +syllapy diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 35d38f992..107fdb775 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -11,6 +11,8 @@ directory is just the launch script + CI test. * `run-qwen2.5-0.5B-fully_async.sh` — single-node, 4-GPU, three-rollout demo with Qwen2.5-0.5B-Instruct on dapo-math-17k. Fast enough to be the CI smoke test for the fully-async path. +* `run-qwen3.5-9B-fully_async.sh` — single-node, 8-GPU, three-rollout demo + with Qwen3.5-9B on dapo-math-17k. The same script doubles as `tests/test_qwen2.5_0.5B_fully_async_short.py` in CI. @@ -56,7 +58,7 @@ work unchanged under fully-async: ``` See `examples/coding_agent_rl/` for a non-trivial example that plugs in a -multi-turn agent this way. +multi-turn agent (Claude Code in a Docker-Proxy sandbox) this way. ## Worker Internals (Very Short) diff --git a/examples/fully_async/run-qwen3-4B-fully_async-npu.sh b/examples/fully_async/run-qwen3-4B-fully_async-npu.sh index 9c05eff36..356d3aa6f 100755 --- a/examples/fully_async/run-qwen3-4B-fully_async-npu.sh +++ b/examples/fully_async/run-qwen3-4B-fully_async-npu.sh @@ -45,7 +45,6 @@ CKPT_ARGS=( --hf-checkpoint "${MODEL_DIR}" --load "${MODEL_DIR}" --ref-load "${MODEL_DIR}" - --megatron-to-hf-mode bridge --save /tmp/vime_fully_async_demo/ --save-interval 9999 ) @@ -106,6 +105,7 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 2 --vllm-gpu-memory-utilization 0.6 ) @@ -119,7 +119,7 @@ MISC_ARGS=( --use-flash-attn ) -ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats +ray start --head --num-gpus 0 --resources '{"NPU": 4}' --node-ip-address 127.0.0.1 --disable-usage-stats # fully-async splits actor / rollout onto disjoint GPUs (no colocation). ACTOR_GPUS=2 diff --git a/examples/fully_async/run-qwen3.5-9B-fully_async.sh b/examples/fully_async/run-qwen3.5-9B-fully_async.sh new file mode 100644 index 000000000..a17259798 --- /dev/null +++ b/examples/fully_async/run-qwen3.5-9B-fully_async.sh @@ -0,0 +1,139 @@ +#!/bin/bash +# Tiny end-to-end fully-async GRPO example using Qwen3.5-9B on the +# dapo-math-17k dataset. Designed to run on a single 8-GPU node in a few +# minutes. +# +# Prerequisites: +# /root/models/Qwen3.5-9B/ (HF checkpoint) +# /root/models/Qwen3.5-9B_torch_dist/ (from tools/convert_hf_to_torch_dist.py) +# /root/datasets/dapo-math-17k/dapo-math-17k.jsonl + +# clean any leftover ray/vllm +pkill -9 -f '[v]llm serve|VLL[M]::' 2>/dev/null || true +sleep 3 +ray stop --force 2>/dev/null || true +pkill -9 ray python 2>/dev/null || true +sleep 3 + +set -ex + +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +HAS_NVLINK=$([ "$NVLINK_COUNT" -gt 0 ] && echo 1 || echo 0) +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/qwen3.5-9B.sh" + +MODEL_DIR=${MODEL_DIR:-/root/models/Qwen3.5-9B} +DATA_PATH=${DATA_PATH:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl} + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_DIR}" + --ref-load "${MODEL_DIR}_torch_dist" + --save /tmp/vime_fully_async_demo/ + --save-interval 9999 +) + +ROLLOUT_ARGS=( + # ↓↓↓ This is the only knob you need to flip to go fully-async ↓↓↓ + --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async + + --prompt-data "${DATA_PATH}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3 + --rollout-batch-size 8 + --n-samples-per-prompt 4 + --rollout-max-response-len 2048 + --rollout-temperature 1 + + --global-batch-size 32 + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 4096 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.7 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +# launch the master node of ray in container +NUM_GPUS=${NUM_GPUS:-8} +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NUM_GPUS}" --disable-usage-stats + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +# fully-async splits actor / rollout onto disjoint GPUs (no colocation). +ACTOR_GPUS=${ACTOR_GPUS:-4} +ROLLOUT_GPUS=${ROLLOUT_GPUS:-$((NUM_GPUS - ACTOR_GPUS))} + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node "${ACTOR_GPUS}" \ + --rollout-num-gpus "${ROLLOUT_GPUS}" \ + ${MODEL_ARGS[@]} \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md index 8e67dd1f8..9805d0a14 100644 --- a/examples/geo3k_vlm/README.md +++ b/examples/geo3k_vlm/README.md @@ -3,21 +3,14 @@ Training VLMs with Megatron on single-turn reasoning task using GRPO on the [GEO3K dataset](https://huggingface.co/datasets/hiyouga/geometry3k). We used processed version [here](https://huggingface.co/datasets/chenhegu/geo3k_imgurl). Supported models: -* Qwen2.5-VL -* Qwen3-VL (Dense and MoE) -* Qwen3.5 (Dense and MoE) +* Qwen3.5-35B-A3B Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prevent severe performance regression in conv3d in torch 2.9 mentioned in https://github.com/pytorch/pytorch/issues/168167. Otherwise, you can reinstall cudnn with: ```bash pip install nvidia-cudnn-cu12==9.16.0.29 ``` -**Important:** We use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to support multimodal models. However, not all Megatron arguments are passed through to Megatron Bridge — you may need to set some manually [here](https://github.com/vllm-project/vime/blob/main/vime/backends/megatron_utils/model_provider.py) (currently only parallelization-related arguments are passed). For example, for Qwen3-VL-30B-A3B you may need to add: -```python -provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff -provider.freeze_language_model = False -provider.freeze_vision_model = False -``` +Qwen3.5 uses vime's native Megatron language model plus the Transformers vision model. Vision parameters retain their HuggingFace names and are loaded and synchronized directly.

GEO3K VLM rollout raw reward @@ -60,45 +53,23 @@ ds.to_parquet("/root/datasets/geo3k_imgurl/train_formatted.parquet") ```bash export WANDB_API_KEY=your_wandb_api_key -# Megatron backend (default -> Qwen3-VL-8B-Instruct + Megatron) -./examples/geo3k_vlm/run_geo3k_vlm.sh - -# With different model -VIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.sh - -# Qwen3.5 ./examples/geo3k_vlm/run_geo3k_qwen35.sh - -# SFT -./examples/geo3k_vlm/run_geo3k_vlm_sft.sh ``` ### Configuration | Environment Variable | Default | Description | |---------------------|---------|-------------| -| `VIME_SCRIPT_MODEL_NAME` | `Qwen3-VL-8B-Instruct` | Model name | | `VIME_SCRIPT_DATASET_NAME` | `chenhegu/geo3k_imgurl` | HuggingFace dataset name | -| `VIME_SCRIPT_NUM_GPUS` | `8` | Number of GPUs used for colocated training and rollout | +| `VIME_SCRIPT_NUM_GPUS` | `8` | Number of GPUs | | `VIME_SCRIPT_EXTERNAL_RAY` | `0` | Use external Ray cluster (`1` to enable) | -### Supported Models - -- `Qwen3-VL-2B-Instruct` -- `Qwen3-VL-4B-Instruct` -- `Qwen3-VL-8B-Instruct` -- `Qwen3-VL-30B-A3B-Instruct` -- `Qwen3-VL-235B-A22B-Instruct` -- `Qwen3-VL-2B-Thinking` -- `Qwen3-VL-4B-Thinking` -- `Qwen3-VL-8B-Thinking` -- `Qwen3-VL-30B-A3B-Thinking` -- `Qwen3-VL-235B-A22B-Thinking` - #### Qwen3.5 Series -We provide an [example](./run_geo3k_qwen35.sh) for Qwen3.5-35B-A3B. To support other Qwen3.5 models, add a model config file in `scripts/models/` and update the model name and config path in the script accordingly. +We provide a native [example](./run_geo3k_qwen35.sh) for Qwen3.5-35B-A3B. It loads the Transformers ViT directly and converts only the Megatron language-model parameters. To support another Qwen3.5 model, add its model config under `scripts/models/` and update the example. + +The native path supports tensor, pipeline, context, sequence, and expert parallelism. Context parallelism follows Megatron's packed THD zigzag layout. -Since Megatron does not currently support packing for GDN, you must set `--qkv-format bshd`, `--micro-batch-size 1`, and remove `--use-dynamic-batch-size`. +For GDN training, use `--micro-batch-size 1` and remove `--use-dynamic-batch-size`. ## Notes @@ -118,7 +89,7 @@ Our initial geo3k-specific verifier produced "format scores" (**0 and 0.9**) ins We fixed this by switching to the default math RM with clean **binary 0/1 rewards**. If you encounter similar precision issues with non-binary rewards, you can change the reward tensor dtype from `torch.float` to `torch.float16` in `vime/ray/rollout.py` (`_post_process_rewards` method) to truncate precision artifacts. ## B200 -On Blackwell (SM100), vllm automatically dispatches the ViT encoder to +On Blackwell (SM100), vLLM automatically dispatches the ViT encoder to FlashAttention 4 (or FA2 fallback) — no manual override is needed ([vllm/v1/attention/backends/fa_utils.py:81](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/fa_utils.py#L81)). If you hit a kernel issue on a specific model, you can force SDPA with diff --git a/examples/geo3k_vlm/run_geo3k_qwen35.sh b/examples/geo3k_vlm/run_geo3k_qwen35.sh index 5026dbda2..e74035190 100644 --- a/examples/geo3k_vlm/run_geo3k_qwen35.sh +++ b/examples/geo3k_vlm/run_geo3k_qwen35.sh @@ -4,16 +4,9 @@ pip install -U transformers -# IMPORTANT: This branch is specially modified for vime's current Megatron -# version and Qwen3.5 from the main Megatron Bridge. Other models are not verified! -# To restore the original Megatron Bridge, run: -# pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation -# TODO: Remove this once Megatron & Megatron Bridge are upgraded upstream. -pip install git+https://github.com/andakai/Megatron-Bridge.git@qwen35 --no-build-isolation - # Configuration TRAIN_BACKEND="megatron" -MODEL_NAME="Qwen3_5-35B-A3B" +MODEL_NAME="Qwen3.5-35B-A3B" DATASET_NAME=${VIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"} NUM_GPUS=${VIME_SCRIPT_NUM_GPUS:-8} DATASET_LOCAL_NAME=$(basename "$DATASET_NAME") @@ -28,7 +21,7 @@ else fi # Cleanup -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 if [ "$USE_EXTERNAL_RAY" = "0" ]; then ray stop --force @@ -68,7 +61,6 @@ fi CKPT_ARGS=( --hf-checkpoint /root/models/${MODEL_NAME} --load /root/models/${MODEL_NAME} - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -119,7 +111,7 @@ VLLM_ARGS=( --rollout-num-gpus-per-engine 8 --vllm-gpu-memory-utilization 0.7 --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + --vllm-cudagraph-capture-sizes 4 8 16 32 $(seq 64 32 1024) # MTP speculative decoding --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' @@ -164,13 +156,11 @@ BACKEND_ARGS=( --attention-softmax-in-fp32 --attention-backend flash - # Packing is not supported for GDN currently - --qkv-format bshd --micro-batch-size 1 ) VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)" -source "${VIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" +source "${VIME_DIR}/scripts/models/qwen3.5-35B-A3B-vl.sh" # Start Ray if not using external Ray if [ "$USE_EXTERNAL_RAY" = "0" ]; then diff --git a/examples/geo3k_vlm/run_geo3k_vlm.sh b/examples/geo3k_vlm/run_geo3k_vlm.sh deleted file mode 100644 index 86f416680..000000000 --- a/examples/geo3k_vlm/run_geo3k_vlm.sh +++ /dev/null @@ -1,218 +0,0 @@ -#!/bin/bash - -# Qwen3 VL RL training on geo3k dataset with colocated vLLM rollout. -# Uses the same GPUs for Megatron training and vLLM rollout. -# Usage: -# VIME_SCRIPT_MODEL_NAME=Qwen3-VL-2B-Instruct ./run_geo3k_vlm.sh - -# Configuration -TRAIN_BACKEND="megatron" -MODEL_NAME=${VIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"} -DATASET_NAME=${VIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"} -NUM_GPUS=${VIME_SCRIPT_NUM_GPUS:-8} -DATASET_LOCAL_NAME=$(basename "$DATASET_NAME") - -# Validate MODEL_NAME -VALID_MODELS=" - Qwen2.5-VL-3B-Instruct - Qwen2.5-VL-7B-Instruct - Qwen2.5-VL-32B-Instruct - Qwen2.5-VL-72B-Instruct - Qwen3-VL-2B-Instruct - Qwen3-VL-4B-Instruct - Qwen3-VL-8B-Instruct - Qwen3-VL-30B-A3B-Instruct - Qwen3-VL-235B-A22B-Instruct - Qwen3-VL-2B-Thinking - Qwen3-VL-4B-Thinking - Qwen3-VL-8B-Thinking - Qwen3-VL-30B-A3B-Thinking - Qwen3-VL-235B-A22B-Thinking -" -if ! echo "$VALID_MODELS" | grep -qw "$MODEL_NAME"; then - echo "Error: MODEL_NAME must be one of: $VALID_MODELS" - exit 1 -fi - -MODEL_NAME_LOWER=$(echo "$MODEL_NAME" | tr '[:upper:]' '[:lower:]') - -# External Ray flag -if [ -z "$VIME_SCRIPT_EXTERNAL_RAY" ] || [ "$VIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then - USE_EXTERNAL_RAY=0 -else - USE_EXTERNAL_RAY=1 -fi - -# Cleanup -sleep 3 -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - ray stop --force - pkill -9 ray -fi -pkill -9 vime -sleep 3 -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - pkill -9 ray -fi -pkill -9 vime -pkill -9 redis - -set -ex - -export PYTHONUNBUFFERED=1 - -# Detect NVLink -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -# Download model and dataset -mkdir -p /root/models /root/datasets -if [ ! -d "/root/models/${MODEL_NAME}" ]; then - hf download Qwen/${MODEL_NAME} --local-dir /root/models/${MODEL_NAME} -fi -if [ ! -d "/root/datasets/${DATASET_LOCAL_NAME}" ]; then - hf download --repo-type dataset ${DATASET_NAME} --local-dir /root/datasets/${DATASET_LOCAL_NAME} -fi - -# Common args -CKPT_ARGS=( - --hf-checkpoint /root/models/${MODEL_NAME} - # qwen3 vl model has rotary base 5000000, set it when applicable - --rotary-base 5000000 -) - -ROLLOUT_ARGS=( - --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train.parquet - --input-key problem - --label-key answer - --apply-chat-template - --rollout-shuffle - --rm-type math - --num-rollout 3000 - --rollout-batch-size 64 - --n-samples-per-prompt 8 - --rollout-max-response-len 4096 - --rollout-temperature 0.8 - --global-batch-size 512 -) - -# required for vlm datasets -MULTIMODAL_KEYS='{"image": "images"}' - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data ${DATASET_LOCAL_NAME} /root/datasets/${DATASET_LOCAL_NAME}/test.parquet - --n-samples-per-eval-prompt 1 - --eval-max-response-len 4096 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-server-concurrency 64 - --vllm-gpu-memory-utilization ${VLLM_GPU_MEMORY_UTILIZATION:-0.9} -) - -# Wandb args (only if WANDB_API_KEY is set) -if [ -n "$WANDB_API_KEY" ]; then - WANDB_ARGS=( - --use-wandb - --wandb-project vime-geo3k-vlm - --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND}-vllm-${NUM_GPUS}gpu-colocate - --wandb-key ${WANDB_API_KEY} - --disable-wandb-random-suffix - ) -else - WANDB_ARGS=() -fi - -MISC_ARGS=( - --colocate -) - -# Backend-specific args -# megatron backend -BACKEND_ARGS=( - --train-backend megatron - --load /root/models/${MODEL_NAME} - --num-gpus-per-node ${NUM_GPUS} - --tensor-model-parallel-size 1 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 4096 - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash - --megatron-to-hf-mode bridge -) - -# get MODEL_ARGS from scripts/models for megatron backend -VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)" -MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g') -# VL models require rotary-base 5000000 -MODEL_ARGS_ROTARY_BASE=5000000 source "${VIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh" - -# Start Ray if not using external Ray -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} - export no_proxy="127.0.0.1,${MASTER_ADDR}" - ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -fi - -# Build runtime env -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node ${NUM_GPUS} \ - --multimodal-keys "${MULTIMODAL_KEYS}" \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${BACKEND_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/examples/geo3k_vlm/run_geo3k_vlm_npu.sh b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh index 911035b22..eb439e590 100644 --- a/examples/geo3k_vlm/run_geo3k_vlm_npu.sh +++ b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh @@ -110,6 +110,7 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 1 --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.8}" --vllm-max-model-len 16384 @@ -130,6 +131,7 @@ else fi BACKEND_ARGS=( + --spec vime_plugins.models.qwen3_vl get_qwen3_vl_model_provider --train-backend megatron --load "$MODEL_ROOT" --tensor-model-parallel-size 4 @@ -148,7 +150,6 @@ BACKEND_ARGS=( --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 --attention-backend flash - --megatron-to-hf-mode bridge ) VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)" @@ -166,7 +167,7 @@ pkill -9 redis || true export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export no_proxy="127.0.0.1,${MASTER_ADDR}" if [ "$USE_EXTERNAL_RAY" = "0" ]; then - ray start --head --node-ip-address "$MASTER_ADDR" --disable-usage-stats \ + ray start --head --num-gpus 0 --resources "{\"NPU\": ${REQUIRED_NPUS}}" --node-ip-address "$MASTER_ADDR" --disable-usage-stats \ --dashboard-host=0.0.0.0 --dashboard-port=8265 fi diff --git a/examples/geo3k_vlm/run_geo3k_vlm_sft.sh b/examples/geo3k_vlm/run_geo3k_vlm_sft.sh deleted file mode 100644 index e976fe35a..000000000 --- a/examples/geo3k_vlm/run_geo3k_vlm_sft.sh +++ /dev/null @@ -1,180 +0,0 @@ -TRAIN_BACKEND="megatron" -MODEL_NAME=${VIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"} -DATASET_NAME=${VIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"} -NUM_GPUS=${VIME_SCRIPT_NUM_GPUS:-8} -DATASET_LOCAL_NAME=$(basename "$DATASET_NAME") - -# Validate MODEL_NAME -VALID_MODELS=" - Qwen2.5-VL-3B-Instruct - Qwen2.5-VL-7B-Instruct - Qwen2.5-VL-32B-Instruct - Qwen2.5-VL-72B-Instruct - Qwen3-VL-2B-Instruct - Qwen3-VL-4B-Instruct - Qwen3-VL-8B-Instruct - Qwen3-VL-2B-Thinking - Qwen3-VL-4B-Thinking - Qwen3-VL-8B-Thinking - Qwen3-VL-30B-A3B-Instruct - Qwen3-VL-235B-A22B-Instruct - Qwen3-VL-30B-A3B-Thinking - Qwen3-VL-235B-A22B-Thinking -" -if ! echo "$VALID_MODELS" | grep -qw "$MODEL_NAME"; then - echo "Error: MODEL_NAME must be one of: $VALID_MODELS" - exit 1 -fi - -MODEL_NAME_LOWER=$(echo "$MODEL_NAME" | tr '[:upper:]' '[:lower:]') - -# External Ray flag -if [ -z "$VIME_SCRIPT_EXTERNAL_RAY" ] || [ "$VIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then - USE_EXTERNAL_RAY=0 -else - USE_EXTERNAL_RAY=1 -fi - -# Cleanup -pkill -9 -f "vllm serve" -sleep 3 -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - ray stop --force - pkill -9 ray -fi -pkill -9 vime -sleep 3 -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - pkill -9 ray -fi -pkill -9 vime -pkill -9 redis - -set -ex - -export PYTHONUNBUFFERED=1 - -# Detect NVLink -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -# Download model and dataset -mkdir -p /root/models /root/datasets -if [ ! -d "/root/models/${MODEL_NAME}" ]; then - hf download Qwen/${MODEL_NAME} --local-dir /root/models/${MODEL_NAME} -fi -if [ ! -d "/root/datasets/${DATASET_LOCAL_NAME}" ]; then - hf download --repo-type dataset ${DATASET_NAME} --local-dir /root/datasets/${DATASET_LOCAL_NAME} -fi - -# Common args -CKPT_ARGS=( - --hf-checkpoint /root/models/${MODEL_NAME} - --load /root/models/${MODEL_NAME} -) - -SFT_ARGS=( - --rollout-function-path vime.rollout.sft_rollout.generate_rollout - --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train_formatted.parquet - --input-key messages - --rollout-shuffle - --num-epoch 3000 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only -) - -# required for vlm datasets -MULTIMODAL_KEYS='{"image": "images"}' - - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-5 - --lr-decay-style cosine - --min-lr 1e-6 - --lr-warmup-fraction 0.1 - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.95 -) - -if [ -n "$WANDB_API_KEY" ]; then - WANDB_ARGS=( - --use-wandb - --wandb-project vime-geo3k-vlm-sft - --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND} - --wandb-key ${WANDB_API_KEY} - --disable-wandb-random-suffix - ) -else - WANDB_ARGS=() -fi - -# Backend-specific args -# megatron backend -BACKEND_ARGS=( - --train-backend megatron - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 4096 - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash - --megatron-to-hf-mode bridge -) - -# get MODEL_ARGS from scripts/models for megatron backend -VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)" -MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g') -# VL models require rotary-base 5000000 -MODEL_ARGS_ROTARY_BASE=5000000 source "${VIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh" - -# Start Ray if not using external Ray -if [ "$USE_EXTERNAL_RAY" = "0" ]; then - export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} - export no_proxy="127.0.0.1,${MASTER_ADDR}" - ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -fi - -# Build runtime env -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train_async.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node ${NUM_GPUS} \ - --multimodal-keys "${MULTIMODAL_KEYS}" \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${SFT_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${BACKEND_ARGS[@]} diff --git a/examples/geo3k_vlm_multi_turn/README.md b/examples/geo3k_vlm_multi_turn/README.md index b162a15b8..3c607a7d6 100644 --- a/examples/geo3k_vlm_multi_turn/README.md +++ b/examples/geo3k_vlm_multi_turn/README.md @@ -27,8 +27,7 @@ The reward model is the default math RM. ```bash # 1) Set environment variable export WANDB_API_KEY=... -export VIME_SCRIPT_MODEL_NAME=Qwen3-VL-2B-Instruct -export VIME_SCRIPT_NUM_GPUS=4 +export VIME_SCRIPT_NUM_GPUS=8 # 2) Download the dataset hf download --repo-type dataset VeraIsHere/geo3k_imgurl_processed --local-dir /root/datasets/geo3k_imgurl_processed @@ -39,7 +38,7 @@ python examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py ``` ## What each file does -- `examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py`: downloads model, sets training/rollout args, and launches the run. +- `examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py`: downloads Qwen3.5-35B-A3B, sets training/rollout args, and launches the run. - `examples/geo3k_vlm_multi_turn/geo3k_vlm_multi_turn_config.yaml`: specifies `max_turns` and `rollout_interaction_env_path` for the multi-turn rollout. - `examples/geo3k_vlm_multi_turn/rollout.py`: custom multi-turn rollout that calls vLLM for token generation (via `/v1/chat/completions/render` → `/inference/v1/generate`), builds loss masks/log_probs, enforces max_turns, and early-stops on max_new_tokens. - `examples/geo3k_vlm_multi_turn/env_geo3k.py`: geo3k tool-calling env that parses {...}, scores math answers, and returns tool feedback per turn. diff --git a/examples/geo3k_vlm_multi_turn/__init__.py b/examples/geo3k_vlm_multi_turn/__init__.py index 526e3ae7c..84920df3d 100644 --- a/examples/geo3k_vlm_multi_turn/__init__.py +++ b/examples/geo3k_vlm_multi_turn/__init__.py @@ -1 +1 @@ -# Multi-turn VLM Sokoban example package +# Multi-turn VLM Geo3K example package diff --git a/examples/geo3k_vlm_multi_turn/rollout.py b/examples/geo3k_vlm_multi_turn/rollout.py index 646186e9d..99d865602 100644 --- a/examples/geo3k_vlm_multi_turn/rollout.py +++ b/examples/geo3k_vlm_multi_turn/rollout.py @@ -106,7 +106,7 @@ def _multimodal_train_inputs_from_features(features: Any) -> dict[str, torch.Ten if not isinstance(encoded_images, list): raise TypeError("vLLM features.kwargs_data.image must be a list") - from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item as vllm_decode + from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item as vllm_decode parts_by_key: dict[str, list[torch.Tensor]] = {} for encoded in encoded_images: @@ -334,28 +334,6 @@ def _remaining_budget(self) -> int | None: return None return self.max_response_budget - self.sample.response_length - def _append_response( - self, - tokens: list[int], - log_probs: list[float] | None = None, - *, - trainable: bool, - meta: dict[str, Any] | None = None, - ) -> None: - # NPU-fork adaptation: this branch's Sample has no append_response_tokens(); maintain the - # token / loss_mask / rollout_log_probs windows directly (equivalent bookkeeping to the rest - # of vime's NPU rollout). Everything else in this file is upstream PR #341 verbatim. - if not tokens: - return - self.sample.tokens.extend(tokens) - self.sample.loss_mask.extend([1 if trainable else 0] * len(tokens)) - self.sample.rollout_log_probs.extend(log_probs if log_probs is not None else [0.0] * len(tokens)) - self.sample.response_length += len(tokens) - if meta and meta.get("routed_experts") is not None: - self.sample.rollout_routed_experts = np.ascontiguousarray( - meta["routed_experts"].astype(np.int32, copy=True) - ) - async def _initialize_prompt(self) -> None: payload: dict[str, Any] = { "model": self.args.hf_checkpoint, @@ -440,7 +418,14 @@ def _append_generated(self, turn: _Turn) -> bool: if eos_token_id is not None and getattr(self.args, "use_rollout_routing_replay", False): raise RuntimeError("Routing replay cannot append an artificial EOS after a stop string") - self._append_response(turn.tokens, turn.log_probs, trainable=True, meta=meta) + self.sample.append_response_tokens( + self.args, + tokens=turn.tokens, + log_probs=turn.log_probs, + trainable=True, + meta_info=meta, + update_terminal_info=False, + ) self.response_tokens.extend(turn.tokens) if eos_token_id is None: return False @@ -452,7 +437,7 @@ def _append_generated(self, turn: _Turn) -> bool: "remaining_budget": remaining, } return True - self._append_response([eos_token_id], trainable=False) + self.sample.append_response_tokens(tokens=[eos_token_id], trainable=False) return False def _advance_environment(self, env: Any, turn: _Turn, turn_index: int) -> bool: @@ -475,7 +460,7 @@ def _advance_environment(self, env: Any, turn: _Turn, turn_index: int) -> bool: ) remaining = self._remaining_budget if remaining is None or len(token_ids) < remaining: - self._append_response(token_ids, trainable=False) + self.sample.append_response_tokens(tokens=token_ids, trainable=False) return False self.sample.status = Sample.Status.TRUNCATED self.sample.metadata["multiturn_truncation"] = { diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py index 390b0b47c..171b27566 100644 --- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py +++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py @@ -3,30 +3,15 @@ import vime.utils.misc as U from vime.utils.external_utils.command_utils import execute_train -MODEL_NAME = os.environ.get("VIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct") -assert MODEL_NAME in { - "Qwen3-VL-2B-Instruct", - "Qwen3-VL-4B-Instruct", - "Qwen3-VL-8B-Instruct", - "Qwen3-VL-2B-Thinking", - "Qwen3-VL-4B-Thinking", - "Qwen3-VL-8B-Thinking", -} +MODEL_NAME = "Qwen3.5-35B-A3B" NUM_GPUS = int(os.environ.get("VIME_SCRIPT_NUM_GPUS", "8")) -EXTERNAL_RAY = int(os.environ.get("VIME_SCRIPT_EXTERNAL_RAY", "0")) DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed" DATA_ROOT = "/root/datasets/geo3k_imgurl_processed" TRAIN_DATA_PATH = os.path.join(DATA_ROOT, "train.parquet") -def get_megatron_model_type(model_name: str) -> str: - model_type = model_name.replace("-Instruct", "").replace("-Thinking", "") - model_type = model_type.replace("Qwen3-VL-", "qwen3-") - return model_type.replace("-2B", "-1.7B") - - def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") @@ -97,43 +82,36 @@ def execute(): cudagraph_sizes = " ".join(map(str, [1, 2, 4, 8] + list(range(16, 257, 8)))) vllm_args = ( - "--rollout-num-gpus-per-engine 1 " - "--vllm-router-policy consistent_hash " + f"--rollout-num-gpus-per-engine {NUM_GPUS} " + "--router-policy consistent_hash " "--vllm-max-model-len 32768 " - "--vllm-gpu-memory-utilization 0.9 " + "--vllm-enable-expert-parallel " + "--vllm-gpu-memory-utilization 0.6 " "--vllm-generation-config vllm " f"--vllm-cudagraph-capture-sizes {cudagraph_sizes} " - # vLLM 0.22.0 needs eager mode for Qwen3-VL logprob parity. This can be - # disabled after vLLM includes https://github.com/vllm-project/vllm/pull/43617. - "--vllm-enforce-eager " "--vllm-logprobs-mode processed_logprobs " ) backend_args = ( "--train-backend megatron " f"--load /root/models/{MODEL_NAME} " - "--tensor-model-parallel-size 1 " + "--tensor-model-parallel-size 2 " "--sequence-parallel " "--pipeline-model-parallel-size 1 " "--context-parallel-size 1 " - "--expert-model-parallel-size 1 " + f"--expert-model-parallel-size {NUM_GPUS} " "--expert-tensor-parallel-size 1 " "--recompute-granularity full " "--recompute-method uniform " "--recompute-num-layers 1 " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 4096 " + "--micro-batch-size 1 " "--attention-dropout 0.0 " "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " - "--megatron-to-hf-mode bridge " ) - megatron_model_type = get_megatron_model_type(MODEL_NAME) - os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000" - misc_args = ( "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_GPUS} " f"--rollout-num-gpus {NUM_GPUS} " "--colocate " ) @@ -153,7 +131,7 @@ def execute(): execute_train( train_args=train_args, num_gpus_per_node=NUM_GPUS, - megatron_model_type=megatron_model_type, + megatron_model_type="qwen3.5-35B-A3B-vl", extra_env_vars=({"WANDB_API_KEY": os.environ["WANDB_API_KEY"]} if os.environ.get("WANDB_API_KEY") else {}), ) diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py index 35124a53a..21be1fe91 100644 --- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py +++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py @@ -89,6 +89,7 @@ def execute(): ) vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.6 " "--vllm-max-model-len 16384 " @@ -117,7 +118,7 @@ def execute(): "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " - "--megatron-to-hf-mode bridge " + "--spec vime_plugins.models.qwen3_vl get_qwen3_vl_model_provider " ) misc_args = ( diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py index 43c04ba4e..eec886467 100644 --- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py +++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py @@ -14,7 +14,7 @@ } EXTERNAL_RAY = int(os.environ.get("VIME_SCRIPT_EXTERNAL_RAY", "0")) -TRAIN_BACKEND = os.environ.get("VIME_SCRIPT_TRAIN_BACKEND", "fsdp").lower() +TRAIN_BACKEND = os.environ.get("VIME_SCRIPT_TRAIN_BACKEND", "megatron").lower() assert TRAIN_BACKEND in {"fsdp", "megatron"} DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed" @@ -95,7 +95,11 @@ def execute(): "--use-precision-aware-optimizer " ) - vllm_args = "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.6 " + vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " + "--rollout-num-gpus-per-engine 1 " + "--vllm-gpu-memory-utilization 0.6 " + ) megatron_args = ( "--train-backend megatron " @@ -118,7 +122,7 @@ def execute(): "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " - "--megatron-to-hf-mode bridge " + "--spec vime_plugins.models.qwen3_vl get_qwen3_vl_model_provider " ) misc_args = ( diff --git a/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh b/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh deleted file mode 100644 index cb5a84db2..000000000 --- a/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh +++ /dev/null @@ -1,15 +0,0 @@ -export VIME_SCRIPT_MODEL_NAME=Qwen3-VL-8B-Instruct -export VIME_SCRIPT_TRAIN_BACKEND=megatron -export PYTORCH_ALLOC_CONF=expandable_segments:True -export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:$PYTHONPATH" -export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 -export CUDA_DEVICE_MAX_CONNECTIONS=1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 -export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 -export HYDRA_FULL_ERROR=1 -export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True -export MASTER_PORT=$(shuf -i 20000-65000 -n 1) # or any free port - - -python examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py diff --git a/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh b/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh deleted file mode 100644 index e67de2e78..000000000 --- a/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh +++ /dev/null @@ -1,15 +0,0 @@ -export VIME_SCRIPT_MODEL_NAME=Qwen3-VL-8B-Instruct -export VIME_SCRIPT_TRAIN_BACKEND=megatron -export PYTORCH_ALLOC_CONF=expandable_segments:True -export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:$PYTHONPATH" -export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 -export CUDA_DEVICE_MAX_CONNECTIONS=1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 -export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 -export HYDRA_FULL_ERROR=1 -export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True -export MASTER_PORT=$(shuf -i 20000-65000 -n 1) # or any free port - - -python examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py diff --git a/examples/mem_agent/README.md b/examples/mem_agent/README.md new file mode 100644 index 000000000..5ee447bec --- /dev/null +++ b/examples/mem_agent/README.md @@ -0,0 +1,115 @@ +# MemAgent on vime + +[MemAgent](https://arxiv.org/abs/2507.02259) (*Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent*) is an **RL-based memory agent** workflow for very long documents. It splits a document into chunks, reads them sequentially, and compresses key information into a **fixed-size memory** (overwrite policy). After all chunks are processed, the model answers using only the problem statement and the memory, with the final answer in `\boxed{}`. Because memory size stays constant, inference scales **linearly** \(O(N)\) with document length—without changing model architecture or positional encodings. + +The paper trains end-to-end with **Multi-Conv DAPO** (this example uses **GRPO**) on multi-turn, context-independent trajectories with verifiable rewards on HotpotQA and RULER. Qwen2.5-7B trained on 32K documents generalizes to million-token QA with near-lossless performance on RULER-HQA. + +This example reproduces that pipeline on vime: HotpotQA multi-turn rollout, GRPO training, and RULER-HQA evaluation, with vLLM as the inference backend. + +The rollout and evaluation implementation was adapted from the `slime-agentic` MemAgent example. + +## Files + +| File | Description | +|------|-------------| +| `rollout.py` | Multi-turn MemAgent rollout + HotpotQA reward | +| `rollout_client.py` | vLLM router client for multi-turn turns | +| `custom_convert.py` | Unroll trajectories for GRPO training | +| `prepare_data.py` | HotpotQA parquet/HF → JSONL | +| `eval_ruler_hqa.py` | RULER-HQA evaluation script | + +## Launch scripts + +Run from **vime repo root** (`cd vime`): + +| Script | Purpose | +|--------|---------| +| `run-qwen3-4b-train.sh` | GRPO training (default 100 steps) | +| `run-eval.sh` | RULER-HQA eval (vLLM serve + `eval_ruler_hqa.py`) | +| `convert-to-hf.sh` | Megatron `iter_*` → HuggingFace | +| `prepare-eval-data.sh` | Check/download `eval_{length}.json` | + +Shared setup lives in `_common.sh` (paths, MemAgent env vars, Ray launch). + +### Quick start + +#### Data download + +Training and evaluation data come from the MemAgent HuggingFace dataset **[BytedTsinghua-SIA/hotpotqa](https://huggingface.co/datasets/BytedTsinghua-SIA/hotpotqa)**. + +**Training set** — download the `train` split and convert to vime JSONL: + +```bash +pip install datasets huggingface_hub pandas pyarrow + +# Optional: use a mirror if huggingface.co is slow +export HF_ENDPOINT=https://hf-mirror.com + +mkdir -p /data/datasets/hotpotqa_slime + +python examples/mem_agent/prepare_data.py \ + --hf-dataset BytedTsinghua-SIA/hotpotqa \ + --hf-split train \ + --output /data/datasets/hotpotqa_slime/train.jsonl +``` + +If you already have a local `hotpotqa_train.parquet`, pass `--input` instead of `--hf-dataset`. + +**Eval set (RULER-HQA)** — `eval_{50,100,200,...}.json` files under `DATA_ROOT` (default `/data/datasets/hotpotqa_hf`): + +```bash +mkdir -p /data/datasets/hotpotqa_hf + +# Download all default lengths (50 … 6400) +bash examples/mem_agent/prepare-eval-data.sh --download + +# Or only the lengths you need +LENGTHS="50 200 800" bash examples/mem_agent/prepare-eval-data.sh --download +``` + +To check which files are present without downloading: + +```bash +LENGTHS="50 200 800" bash examples/mem_agent/prepare-eval-data.sh +``` + +#### Run pipeline + +```bash +cd vime + +# 1. Training (100 steps by default; set TRAIN_DATA if you used a different path) +bash examples/mem_agent/run-qwen3-4b-train.sh + +# 2. Eval (after convert to HF) +CONVERT=1 SINGLE_ITER=iter_0000099 bash examples/mem_agent/run-eval.sh + +# Baseline (untrained Qwen3-4B) +MODEL_PATH=/data/models/Qwen3-4B SAVE_FILE=Qwen3-4B-base bash examples/mem_agent/run-eval.sh +``` + +### Environment variables + +Paths (override for your cluster): + +- `HF_CKPT`, `TORCH_DIST`, `TRAIN_DATA`, `SAVE_PATH`, `DATA_ROOT` + +Training: + +- `NUM_ROLLOUT`, `MEM_CHUNK_TOKENS`, `MEM_MAX_MEMORY`, `MEM_MAX_FINAL`, `MEM_MAX_CHUNKS` + +Eval: + +- `MODEL_PATH`, `LENGTH` (e.g. `"50 200 800"`), `CONVERT`, `SINGLE_ITER`, `SAVE_FILE` + +### Example paths + +```bash +export HF_CKPT=/data/models/Qwen3-4B +export TORCH_DIST=/data/models/Qwen3-4B_torch_dist +export TRAIN_DATA=/data/datasets/hotpotqa_slime/train.jsonl +export SAVE_PATH=/data/models/MemAgent_Qwen3-4B-RL +export DATA_ROOT=/data/datasets/hotpotqa_hf + +bash examples/mem_agent/run-qwen3-4b-train.sh +``` diff --git a/examples/mem_agent/_common.sh b/examples/mem_agent/_common.sh new file mode 100644 index 000000000..d08bf2ca6 --- /dev/null +++ b/examples/mem_agent/_common.sh @@ -0,0 +1,203 @@ +#!/usr/bin/env bash +# Shared setup for MemAgent example scripts (sourced, not executed directly). +set -euo pipefail + +MEM_AGENT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_ROOT="$(cd -- "${MEM_AGENT_DIR}/../.." &>/dev/null && pwd)" + +export VIME_ROOT +export PYTHONBUFFERED="${PYTHONBUFFERED:-16}" +export WANDB_MODE="${WANDB_MODE:-disabled}" +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +export no_proxy="127.0.0.1,${MASTER_ADDR},localhost" +export NO_PROXY="${no_proxy}" + +# Override via env for your cluster (H800 example in README). +export HF_CKPT="${HF_CKPT:-/data/models/Qwen3-4B}" +export TORCH_DIST="${TORCH_DIST:-/data/models/Qwen3-4B_torch_dist}" +export TRAIN_DATA="${TRAIN_DATA:-/data/datasets/hotpotqa_slime/train.jsonl}" +export SAVE_PATH="${SAVE_PATH:-/data/models/MemAgent_Qwen3-4B-RL}" +export DATA_ROOT="${DATA_ROOT:-/data/datasets/hotpotqa_hf}" +export ORIGIN_HF_DIR="${ORIGIN_HF_DIR:-${HF_CKPT}}" + +export MEM_CHUNK_TOKENS="${MEM_CHUNK_TOKENS:-2048}" +export MEM_MAX_MEMORY="${MEM_MAX_MEMORY:-1024}" +export MEM_MAX_FINAL="${MEM_MAX_FINAL:-256}" +export MEM_MAX_CHUNKS="${MEM_MAX_CHUNKS:-64}" + +mem_agent_detect_nvlink() { + local count + count=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) + export NCCL_NVLS_ENABLE=$([[ "${count}" -gt 0 ]] && echo 1 || echo 0) +} + +mem_agent_detect_gpus() { + local detected=0 + if command -v nvidia-smi >/dev/null 2>&1; then + detected=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ') + fi + export NUM_GPUS="${NUM_GPUS:-${detected:-8}}" + if [[ -z "${NUM_GPUS}" || "${NUM_GPUS}" -le 0 ]]; then + export NUM_GPUS=8 + fi +} + +mem_agent_cleanup() { + pkill -9 -f '[v]llm serve|VLLM::' 2>/dev/null || true + sleep 2 + ray stop --force 2>/dev/null || true + pkill -9 -f '[r]ay::' 2>/dev/null || true + pkill -9 -f '[t]rain.py' 2>/dev/null || true + pkill -9 redis 2>/dev/null || true + rm -rf /tmp/ray /tmp/ray_session_* "${HOME}/.ray" 2>/dev/null || true + sleep 2 +} + +mem_agent_rollout_args() { + ROLLOUT_ARGS=( + --custom-generate-function-path examples.mem_agent.rollout.generate + --custom-rm-path examples.mem_agent.rollout.reward_func + --custom-convert-samples-to-train-data-path examples.mem_agent.custom_convert.custom_convert + --prompt-data "${TRAIN_DATA}" + --input-key prompt + --label-key label + --rollout-shuffle + --reward-key score + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-max-response-len 8192 + --rollout-max-context-len 32768 + --rollout-temperature 1.0 + --rollout-top-p 1.0 + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --balance-data + --rollout-function-path vime.rollout.vllm_rollout.generate_rollout + ) +} + +mem_agent_train_args() { + CKPT_ARGS=( + --hf-checkpoint "${HF_CKPT}" + --ref-load "${TORCH_DIST}" + ) + if [[ -n "${SAVE_PATH:-}" ]]; then + CKPT_ARGS+=(--save "${SAVE_PATH}") + if [[ -n "${SAVE_INTERVAL:-}" ]]; then + CKPT_ARGS+=(--save-interval "${SAVE_INTERVAL}") + fi + fi + + EVAL_ARGS=(--skip-eval-before-train) + + PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 + ) + + GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.001 + --kl-loss-type low_var_kl + --eps-clip 0.2 + --eps-clip-high 0.3 + ) + + OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + ) + + VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.7 + --vllm-max-model-len 32768 + --router-policy consistent_hash + ) + + MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --actor-num-nodes 1 + --actor-num-gpus-per-node "${NUM_GPUS}" + --colocate + ) +} + +mem_agent_launch_train() { + cd "${VIME_ROOT}" + source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" + mem_agent_rollout_args + mem_agent_train_args + + export RAY_DISABLE_DOCKER_CPU_WARNING=1 + export PYTHONPATH="/root/Megatron-LM:${VIME_ROOT}:${PYTHONPATH:-}" + + if [[ "${RUN_TRAIN_DIRECT:-0}" == "1" && -n "${RUN_TRAIN_DIRECT_PY:-}" && -f "${RUN_TRAIN_DIRECT_PY}" ]]; then + export VIME_ROOT NCCL_NVLS_ENABLE MEM_CHUNK_TOKENS MEM_MAX_MEMORY MEM_MAX_FINAL MEM_MAX_CHUNKS + python3 "${RUN_TRAIN_DIRECT_PY}" \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" + return + fi + + ray start --head --node-ip-address "${MASTER_ADDR}" \ + --num-gpus "${NUM_GPUS}" --disable-usage-stats \ + --dashboard-host=0.0.0.0 --dashboard-port=8265 + + local runtime_env + runtime_env="{ + \"env_vars\": { + \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM/\", + \"VIME_ROOT\": \"${VIME_ROOT}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${NCCL_NVLS_ENABLE}\", + \"PYTHONBUFFERED\": \"16\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"no_proxy\": \"${no_proxy}\", + \"NO_PROXY\": \"${NO_PROXY}\", + \"MEM_CHUNK_TOKENS\": \"${MEM_CHUNK_TOKENS}\", + \"MEM_MAX_MEMORY\": \"${MEM_MAX_MEMORY}\", + \"MEM_MAX_FINAL\": \"${MEM_MAX_FINAL}\", + \"MEM_MAX_CHUNKS\": \"${MEM_MAX_CHUNKS}\" + } + }" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${runtime_env}" \ + -- python3 train.py \ + --train-backend megatron \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" +} diff --git a/examples/mem_agent/convert-to-hf.sh b/examples/mem_agent/convert-to-hf.sh new file mode 100644 index 000000000..c1185afbf --- /dev/null +++ b/examples/mem_agent/convert-to-hf.sh @@ -0,0 +1,61 @@ +#!/usr/bin/env bash +# Convert MemAgent Megatron checkpoints to HuggingFace format. +# +# Usage: +# bash examples/mem_agent/convert-to-hf.sh +# CHECKPOINT_DIR=/path/to/ckpt SINGLE_ITER=iter_0000199 bash examples/mem_agent/convert-to-hf.sh +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# shellcheck source=_common.sh +source "${SCRIPT_DIR}/_common.sh" + +CHECKPOINT_DIR="${CHECKPOINT_DIR:-${SAVE_PATH}}" +OUTPUT_BASE="${OUTPUT_BASE:-${CHECKPOINT_DIR}-HF}" +MEGATRON_LM_DIR="${MEGATRON_LM_DIR:-/root/Megatron-LM}" +CONVERT_SCRIPT="${CONVERT_SCRIPT:-${VIME_ROOT}/tools/convert_torch_dist_to_hf.py}" + +mkdir -p "${OUTPUT_BASE}" +export PYTHONPATH="${MEGATRON_LM_DIR}:${VIME_ROOT}:${PYTHONPATH:-}" + +if [[ -n "${SINGLE_ITER:-}" ]]; then + ITERS=("${CHECKPOINT_DIR}/${SINGLE_ITER}") +else + mapfile -t ITERS < <(ls -d "${CHECKPOINT_DIR}"/iter_* 2>/dev/null | sort) +fi + +if [[ ${#ITERS[@]} -eq 0 ]]; then + echo "ERROR: no iter_* checkpoints in ${CHECKPOINT_DIR}" + exit 1 +fi + +echo "Converting ${#ITERS[@]} checkpoint(s) -> ${OUTPUT_BASE}" +FAILED=() + +for iter_path in "${ITERS[@]}"; do + iter_name="$(basename "${iter_path}")" + output_dir="${OUTPUT_BASE}/${iter_name}" + + if [[ -d "${output_dir}" && -f "${output_dir}/config.json" ]]; then + echo "[SKIP] ${iter_name} already at ${output_dir}" + continue + fi + + echo "[CONVERT] ${iter_name} -> ${output_dir}" + if python3 "${CONVERT_SCRIPT}" \ + --input-dir "${iter_path}" \ + --output-dir "${output_dir}" \ + --origin-hf-dir "${ORIGIN_HF_DIR}"; then + echo "[DONE] ${iter_name}" + else + echo "[FAILED] ${iter_name}" + FAILED+=("${iter_name}") + fi +done + +echo "===== Summary: total=${#ITERS[@]} failed=${#FAILED[@]} =====" +if [[ ${#FAILED[@]} -gt 0 ]]; then + printf ' %s\n' "${FAILED[@]}" + exit 1 +fi +echo "Output: ${OUTPUT_BASE}" diff --git a/examples/mem_agent/custom_convert.py b/examples/mem_agent/custom_convert.py new file mode 100644 index 000000000..248b2cdd0 --- /dev/null +++ b/examples/mem_agent/custom_convert.py @@ -0,0 +1,137 @@ +"""MemAgent multi-turn custom convert — unroll turns into independent training samples.""" + +from __future__ import annotations + +import logging + +import torch + +logger = logging.getLogger(__name__) + + +def _sample_has_rollout_log_probs(sample) -> bool: + meta = sample.train_metadata + if meta and "turns" in meta: + return any(t.get("rollout_log_probs") is not None for t in meta["turns"]) + return sample.rollout_log_probs is not None + + +def custom_convert(args, samples): + raw_rewards = [s.get_reward_value(args) for s in samples] + rewards_tensor = torch.tensor(raw_rewards, dtype=torch.float) + + if getattr(args, "advantage_estimator", None) in ["grpo", "gspo", "reinforce_plus_plus_baseline"] and getattr( + args, "rewards_normalization", False + ): + n = getattr(args, "n_samples_per_prompt", 1) + if rewards_tensor.shape[-1] == n * getattr(args, "rollout_batch_size", 1): + rewards_tensor = rewards_tensor.reshape(-1, n) + else: + rewards_tensor = rewards_tensor.view(-1, rewards_tensor.shape[-1]) + + mean = rewards_tensor.mean(dim=-1, keepdim=True) + rewards_tensor = rewards_tensor - mean + + if getattr(args, "advantage_estimator", None) in ["grpo", "gspo"] and getattr( + args, "grpo_std_normalization", False + ): + std = rewards_tensor.std(dim=-1, keepdim=True) + rewards_tensor = rewards_tensor / (std + 1e-6) + + normalized_rewards = rewards_tensor.flatten().tolist() + + tokens_list = [] + response_lengths = [] + loss_masks = [] + rewards = [] + raw_reward_list = [] + truncated_list = [] + sample_indices = [] + rollout_log_probs_list = [] + has_rollout_log_probs = any( + _sample_has_rollout_log_probs(s) for s in samples if s.status not in (s.Status.FAILED, s.Status.ABORTED) + ) + + for i, sample in enumerate(samples): + if sample.status in (sample.Status.FAILED, sample.Status.ABORTED): + continue + + meta = sample.train_metadata + if meta is None or "turns" not in meta: + if not sample.tokens: + continue + tokens_list.append(sample.tokens) + response_lengths.append(sample.response_length) + lm = sample.loss_mask if sample.loss_mask is not None else [1] * sample.response_length + if sample.remove_sample: + lm = [0] * sample.response_length + loss_masks.append(lm) + rewards.append(normalized_rewards[i]) + raw_reward_list.append(raw_rewards[i]) + truncated_list.append(1 if sample.status == sample.Status.TRUNCATED else 0) + sample_indices.append(sample.index) + if has_rollout_log_probs: + lp = sample.rollout_log_probs + if lp is None: + lp = [0.0] * sample.response_length + rollout_log_probs_list.append(lp) + continue + + turns = meta["turns"] + if not turns: + continue + norm_reward = normalized_rewards[i] / len(turns) + is_truncated = 1 if sample.status == sample.Status.TRUNCATED else 0 + + for turn in turns: + tokens_list.append(turn["tokens"]) + response_lengths.append(turn["response_length"]) + lm = list(turn["loss_mask"]) + if sample.remove_sample: + lm = [0] * turn["response_length"] + loss_masks.append(lm) + rewards.append(norm_reward) + raw_reward_list.append(raw_rewards[i]) + truncated_list.append(is_truncated) + sample_indices.append(sample.index) + if has_rollout_log_probs: + lp = turn.get("rollout_log_probs") + if lp is None: + lp = [0.0] * turn["response_length"] + rollout_log_probs_list.append(lp) + + gbs = args.global_batch_size + total = len(tokens_list) + trim_to = (total // gbs) * gbs + if trim_to == 0: + trim_to = total + if trim_to < total: + logger.info( + "custom_convert: trimming expanded samples from %d to %d (global_batch_size=%d)", + total, + trim_to, + gbs, + ) + tokens_list = tokens_list[:trim_to] + response_lengths = response_lengths[:trim_to] + loss_masks = loss_masks[:trim_to] + rewards = rewards[:trim_to] + raw_reward_list = raw_reward_list[:trim_to] + truncated_list = truncated_list[:trim_to] + sample_indices = sample_indices[:trim_to] + if has_rollout_log_probs: + rollout_log_probs_list = rollout_log_probs_list[:trim_to] + + train_data = { + "tokens": tokens_list, + "response_lengths": response_lengths, + "loss_masks": loss_masks, + "rewards": rewards, + "raw_reward": raw_reward_list, + "truncated": truncated_list, + "sample_indices": sample_indices, + } + if has_rollout_log_probs: + train_data["rollout_log_probs"] = rollout_log_probs_list + + return train_data diff --git a/examples/mem_agent/eval_ruler_hqa.py b/examples/mem_agent/eval_ruler_hqa.py new file mode 100644 index 000000000..9b53eeb4d --- /dev/null +++ b/examples/mem_agent/eval_ruler_hqa.py @@ -0,0 +1,542 @@ +""" +eval_ruler_hqa.py — MemAgent HotpotQA evaluation (vime / vLLM version) +====================================================================== +Core logic aligned with slime-agentic eval_ruler_hqa.py and MemAgent ruler_hqa.py. +Inference uses vLLM OpenAI-compatible ``/v1/chat/completions``. + +Evaluation data: MemAgent-format eval_{length}.json + Fields: input, answers, context, num_docs + +Metrics: F1, EM, sub_EM + +Usage (vLLM server must be running, e.g. via examples/mem_agent/run-eval.sh): + python eval_ruler_hqa.py \\ + --model /path/to/hf_ckpt \\ + --tokenizer /path/to/hf_ckpt \\ + --length 200 \\ + --data-root /data/hotpotqa \\ + --save-dir results/ruler_hqa_200 \\ + --save-file iter_0000200 + +Environment variables: + VLLM_SERVE_HOST / SERVE_HOST vLLM server host (default: 127.0.0.1) + VLLM_SERVE_PORT / SERVE_PORT vLLM server port (default: 8000) + MEM_CHUNK_TOKENS tokens per chunk (default: 2048) + MEM_MAX_MEMORY max tokens for memory update (default: 1024) + MEM_MAX_FINAL max tokens for final answer (default: 256) + MEM_MAX_CHUNKS max number of chunks (default: 512) + DATAROOT directory with eval_*.json +""" + +from __future__ import annotations + +import argparse +import asyncio +import copy +import json +import os +import re +import string +from collections import Counter +from pathlib import Path + +import aiohttp +from tqdm import tqdm +from transformers import AutoTokenizer + +SERVE_HOST = os.getenv("VLLM_SERVE_HOST", os.getenv("SERVE_HOST", "127.0.0.1")) +SERVE_PORT = os.getenv("VLLM_SERVE_PORT", os.getenv("SERVE_PORT", "8000")) +CHUNK_TOKENS = int(os.getenv("MEM_CHUNK_TOKENS", "2048")) +MAX_MEMORY_TOKS = int(os.getenv("MEM_MAX_MEMORY", "1024")) +MAX_FINAL_TOKS = int(os.getenv("MEM_MAX_FINAL", "256")) +MAX_CHUNKS = int(os.getenv("MEM_MAX_CHUNKS", "512")) +MAX_CTX_TOKENS = int(os.getenv("MEM_MAX_CTX_TOKENS", str(10**12))) +BASE_URL = f"http://{SERVE_HOST}:{SERVE_PORT}/v1" +API_KEY = os.getenv("SERVE_API_KEY", "EMPTY") +DATAROOT = os.getenv("DATAROOT", "/data/hotpotqa") + +_MEMORY_TEMPLATE = """You are presented with a problem, a section of an article that may contain the answer to the problem, and a previous memory. Please read the provided section carefully and update the memory with the new information that helps to answer the problem. Be sure to retain all relevant details from the previous memory while adding any new, useful information. + + +{prompt} + + + +{memory} + + +

+{chunk} +
+ +Updated memory: +""" + +_FINAL_TEMPLATE = """You are presented with a problem and a previous memory. Please answer the problem based on the previous memory and put the answer in \\boxed{{}}. + + +{prompt} + + + +{memory} + + +Your answer: +""" + +_NO_MEMORY = "No previous memory" +_STOP_TOKEN_STRINGS = ["<|im_end|>", "<|endoftext|>"] + + +def _strip_stop_tokens(text: str) -> str: + for tok in _STOP_TOKEN_STRINGS: + text = text.replace(tok, "") + return text.strip() + + +def _last_boxed_only_string(s: str) -> str | None: + if "\\boxed " in s: + return "\\boxed " + s.split("\\boxed ")[-1].split("$")[0] + idx = s.rfind("\\boxed") + if idx < 0: + idx = s.rfind("\\fbox") + if idx < 0: + return None + i, right_brace_idx, opens = idx, None, 0 + while i < len(s): + if s[i] == "{": + opens += 1 + if s[i] == "}": + opens -= 1 + if opens == 0: + right_brace_idx = i + break + i += 1 + return s[idx : right_brace_idx + 1] if right_brace_idx is not None else None + + +def _remove_boxed(s: str) -> str: + if s.startswith("\\boxed "): + return s[len("\\boxed ") :] + left = "\\boxed{" + assert s.startswith(left) and s.endswith("}") + return s[len(left) : -1] + + +def _extract_boxed(text: str) -> str: + s = _last_boxed_only_string(text) + if s is None: + return "" + try: + return _remove_boxed(s).strip() + except Exception: + return "" + + +def _normalize_answer(s: str) -> str: + def remove_articles(t): + return re.sub(r"\b(a|an|the)\b", " ", t) + + def white_space_fix(t): + return " ".join(t.split()) + + def remove_punc(t): + exclude = set(string.punctuation) + return "".join(ch for ch in t if ch not in exclude) + + return white_space_fix(remove_articles(remove_punc(s.lower()))) + + +def _f1_score(prediction: str, ground_truth: str) -> tuple[float, float, float]: + p = _normalize_answer(prediction).split() + g = _normalize_answer(ground_truth).split() + if not p or not g: + return 0.0, 0.0, 0.0 + common = Counter(p) & Counter(g) + num_same = sum(common.values()) + if num_same == 0: + return 0.0, 0.0, 0.0 + prec = num_same / len(p) + recall = num_same / len(g) + f1 = 2 * prec * recall / (prec + recall) + return f1, prec, recall + + +def _exact_match(prediction: str, ground_truth: str) -> float: + p = _normalize_answer(prediction) + g = _normalize_answer(ground_truth) + if p in ("yes", "no", "noanswer") and p != g: + return 0.0 + if g in ("yes", "no", "noanswer") and p != g: + return 0.0 + return float(p == g) + + +def _sub_exact_match(prediction: str, ground_truth: str) -> float: + p = _normalize_answer(prediction) + g = _normalize_answer(ground_truth) + return float((g in p) or (p in g)) + + +def _agg_metrics(records: list[dict]) -> dict: + keys = ("judge_f1", "judge_em", "judge_sub_em") + totals = {k: 0.0 for k in keys} + n = len(records) + for r in records: + for k in keys: + totals[k] += r[k] + return { + "f1": round(totals["judge_f1"] / n, 4) if n else 0.0, + "em": round(totals["judge_em"] / n, 4) if n else 0.0, + "sub_em": round(totals["judge_sub_em"] / n, 4) if n else 0.0, + "total": n, + } + + +async def _chat_once( + session: aiohttp.ClientSession, + model: str, + messages: list[dict], + temperature: float, + top_p: float, + max_tokens: int, +) -> str: + payload = dict( + model=model, + messages=messages, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + ) + async with session.post( + f"{BASE_URL}/chat/completions", + headers={"Authorization": f"Bearer {API_KEY}"}, + json=payload, + ) as resp: + if resp.status != 200: + body = await resp.text() + raise RuntimeError(f"HTTP {resp.status}: {body[:300]}") + data = await resp.json() + return data["choices"][0]["message"]["content"] + + +async def _recurrent_infer( + item: dict, + model: str, + tokenizer, + temperature: float, + top_p: float, + sem: asyncio.Semaphore, +) -> str: + question = item["input"].strip() + context = item["context"].strip() + + ctx_ids = tokenizer.encode(context, add_special_tokens=False) + if len(ctx_ids) > MAX_CTX_TOKENS: + half = MAX_CTX_TOKENS // 2 + ctx_ids = ctx_ids[:half] + ctx_ids[-half:] + + chunks = [ctx_ids[i : i + CHUNK_TOKENS] for i in range(0, len(ctx_ids), CHUNK_TOKENS)][:MAX_CHUNKS] + + async with sem: + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=86400)) as session: + memory = _NO_MEMORY + for chunk_ids in chunks: + chunk_text = tokenizer.decode(chunk_ids, skip_special_tokens=True) + msg = _MEMORY_TEMPLATE.format(prompt=question, memory=memory, chunk=chunk_text) + raw = await _chat_once( + session, + model, + [{"role": "user", "content": msg}], + temperature, + top_p, + MAX_MEMORY_TOKS, + ) + memory = _strip_stop_tokens(raw) or memory + + final_msg = _FINAL_TEMPLATE.format(prompt=question, memory=memory) + response = await _chat_once( + session, + model, + [{"role": "user", "content": final_msg}], + temperature, + top_p, + MAX_FINAL_TOKS, + ) + return response.strip() + + +async def _openai_infer( + item: dict, + model: str, + tokenizer, + temperature: float, + top_p: float, + max_input_len: int, + max_output_len: int, + sem: asyncio.Semaphore, +) -> str: + question = item["input"].strip() + context = item["context"].strip() + prompt = f"{context}\n\n" f"Question: {question}\n" "Please put the answer in \\boxed{}.\n\n" "Answer:" + ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(ids) > max_input_len: + prompt = tokenizer.decode(ids[:max_input_len], skip_special_tokens=True) + + async with sem: + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=86400)) as session: + response = await _chat_once( + session, + model, + [{"role": "user", "content": prompt}], + temperature, + top_p, + max_output_len, + ) + return response.strip() + + +async def run_eval(data: list[dict], args: argparse.Namespace, tokenizer) -> None: + out_path = Path(args.save_dir) / f"{args.save_file}.jsonl" + out_path.parent.mkdir(parents=True, exist_ok=True) + + cached_ids: set[int] = set() + if out_path.exists() and not args.force: + with open(out_path, encoding="utf-8") as f: + for line in f: + try: + cached_ids.add(json.loads(line)["_id"]) + except Exception: + pass + + todo = [item for item in data if item["_id"] not in cached_ids] + print(f"Data: {len(data)} Cached: {len(cached_ids)} " f"Todo: {len(todo)} Concurrency: {args.n_proc}") + if not todo: + print("All done (cached).") + _print_existing_stats(out_path) + return + + sem = asyncio.Semaphore(args.n_proc) + + async def process_one(item: dict) -> dict | None: + try: + if args.api == "recurrent": + response = await _recurrent_infer(item, args.model, tokenizer, args.temperature, args.top_p, sem) + else: + response = await _openai_infer( + item, + args.model, + tokenizer, + args.temperature, + args.top_p, + args.max_input_len, + args.max_output_len, + sem, + ) + except Exception: + import traceback + + traceback.print_exc() + return None + + gold = item["answers"][0] if item.get("answers") else "" + pred = _extract_boxed(response[-300:]) or "" + + result = { + "_id": item["_id"], + "answer": gold, + "pred": pred, + "judge_f1": _f1_score(pred, gold)[0], + "judge_em": _exact_match(pred, gold), + "judge_sub_em": _sub_exact_match(pred, gold), + "response": response, + } + for k, v in item.items(): + if k not in ("context", "response") and k not in result: + result[k] = v + return result + + tasks = [process_one(item) for item in todo] + records: list[dict] = [] + n_err = 0 + first_shown = False + fout = open(out_path, "a", encoding="utf-8") + pbar = tqdm(total=len(tasks), desc=f"ruler_hqa[{args.length}]", dynamic_ncols=True) + PRINT_INTERVAL = 10 + + for coro in asyncio.as_completed(tasks): + result = await coro + pbar.update(1) + if result is None: + n_err += 1 + pbar.set_postfix(err=n_err, done=len(records)) + continue + + fout.write(json.dumps(result, ensure_ascii=False) + "\n") + fout.flush() + records.append(result) + + n = len(records) + rolling_f1 = sum(r["judge_f1"] for r in records) / n + rolling_sub_em = sum(r["judge_sub_em"] for r in records) / n + pbar.set_postfix( + done=n, + err=n_err, + f1=f"{rolling_f1 * 100:.1f}", + sub_em=f"{rolling_sub_em * 100:.1f}", + ) + + if not first_shown: + first_shown = True + _print_sample(result) + + if n % PRINT_INTERVAL == 0: + tqdm.write( + f"[interim {n}/{len(tasks)}] " + f"F1={rolling_f1 * 100:.2f} " + f"sub_EM={rolling_sub_em * 100:.2f} " + f"err={n_err}" + ) + + pbar.close() + fout.close() + + all_records = records[:] + if cached_ids: + with open(out_path, encoding="utf-8") as f: + all_records = [json.loads(line) for line in f if line.strip()] + + stats = _agg_metrics(all_records) + print(f"\n=== ruler_hqa [n_docs={args.length}] " f"total={stats['total']} errors={n_err} ===") + for k in ("f1", "em", "sub_em"): + print(f" {k}: {round(stats[k] * 100, 2)}") + + +def _print_sample(result: dict) -> None: + sep = "=" * 40 + print(f"\n{sep} Sample {result['_id']} {sep}") + print(f"[response] {result['response'][:500]}") + print(f"[pred] {result['pred']}") + print(f"[answer] {result['answer']}") + print(f"[sub_em] {result['judge_sub_em']}") + print(sep) + + +def _print_existing_stats(out_path: Path) -> None: + records = [] + with open(out_path, encoding="utf-8") as f: + for line in f: + try: + records.append(json.loads(line)) + except Exception: + pass + if not records: + return + stats = _agg_metrics(records) + print(f"Existing results ({stats['total']} samples):") + for k in ("f1", "em", "sub_em"): + print(f" {k}: {round(stats[k] * 100, 2)}") + + +def load_data(data_root: str, length: int) -> list[dict]: + candidates = [ + Path(data_root) / f"eval_{length}.json", + Path(data_root) / f"eval_{length}.jsonl", + ] + data_path = next((p for p in candidates if p.exists()), None) + if data_path is None: + raise FileNotFoundError( + f"Cannot find eval data for length={length} in {data_root!r}. " f"Tried: {[str(p) for p in candidates]}" + ) + + suffix = data_path.suffix.lower() + if suffix == ".json": + with open(data_path, encoding="utf-8") as f: + raw = json.load(f) + if isinstance(raw, dict): + raw = list(raw.values()) + else: + with open(data_path, encoding="utf-8") as f: + raw = [json.loads(line) for line in f if line.strip()] + + data = [] + for idx, item in enumerate(raw): + if "input" in item: + item = dict(item) + item.setdefault("_id", idx) + data.append(item) + elif "prompt" in item: + meta = item.get("metadata") or {} + data.append( + { + "_id": idx, + "input": item["prompt"], + "answers": meta.get("ground_truth", [item.get("label", "")]), + "context": meta.get("context", ""), + "num_docs": meta.get("num_docs", 0), + } + ) + else: + print(f"[warn] skipping row {idx}: unrecognized format (keys={list(item.keys())[:5]})") + + return data + + +def main() -> None: + parser = argparse.ArgumentParser(description="MemAgent HotpotQA eval — vime/vLLM version") + parser.add_argument( + "--length", + type=int, + default=200, + choices=[50, 100, 200, 400, 800, 1600, 3200, 6400, 12800, 25600], + help="Number of distractive documents (controls context length)", + ) + parser.add_argument( + "--data-root", + default=DATAROOT, + help="Directory containing eval_{length}.json files", + ) + parser.add_argument("--save-dir", "-s", default="results/ruler_hqa", help="Output directory") + parser.add_argument("--save-file", "-f", default="model", help="Output filename stem") + parser.add_argument("--model", "-m", required=True, help="Model id registered in vLLM server") + parser.add_argument("--tokenizer", "-t", required=True, help="HuggingFace tokenizer path (for chunking)") + parser.add_argument( + "--api", + default="recurrent", + choices=["recurrent", "openai"], + help="recurrent: chunk-by-chunk memory update (default); openai: single-turn baseline", + ) + parser.add_argument("--n-proc", "-n", type=int, default=32, help="Max concurrent requests") + parser.add_argument("--temperature", type=float, default=0.7) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--max-input-len", type=int, default=120000) + parser.add_argument("--max-output-len", type=int, default=10000) + parser.add_argument("--sampling", type=int, default=1) + parser.add_argument("--force", action="store_true", help="Ignore cache and re-evaluate") + args = parser.parse_args() + + print(f"[config] vLLM={SERVE_HOST}:{SERVE_PORT} api={args.api}") + print( + f"[config] CHUNK_TOKENS={CHUNK_TOKENS} MAX_MEMORY={MAX_MEMORY_TOKS} " + f"MAX_FINAL={MAX_FINAL_TOKS} MAX_CHUNKS={MAX_CHUNKS}" + ) + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + data = load_data(args.data_root, args.length) + + if args.sampling > 1: + base = data[:] + data = [] + for s in range(args.sampling): + for item in base: + new_item = copy.deepcopy(item) + new_item["_id"] = item["_id"] * args.sampling + s + data.append(new_item) + + print(f"[data] {len(data)} samples (n_docs={args.length})") + asyncio.run(run_eval(data, args, tokenizer)) + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/prepare-eval-data.sh b/examples/mem_agent/prepare-eval-data.sh new file mode 100644 index 000000000..110a73065 --- /dev/null +++ b/examples/mem_agent/prepare-eval-data.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Prepare RULER-HQA eval JSON (eval_{50,100,200,...}.json). +# +# Usage: +# bash examples/mem_agent/prepare-eval-data.sh +# LENGTHS="50 200 800" bash examples/mem_agent/prepare-eval-data.sh +# bash examples/mem_agent/prepare-eval-data.sh --download +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# shellcheck source=_common.sh +source "${SCRIPT_DIR}/_common.sh" + +DATA_DIR="${DATA_DIR:-${DATA_ROOT}}" +HF_DATASET="${HF_DATASET:-BytedTsinghua-SIA/hotpotqa}" +LENGTHS="${LENGTHS:-50 100 200 400 800 1600 3200 6400}" + +mkdir -p "${DATA_DIR}" +export PYTHONPATH="/root/Megatron-LM:${VIME_ROOT}:${PYTHONPATH:-}" + +download_one() { + local length="$1" + local fname="eval_${length}.json" + local dest="${DATA_DIR}/${fname}" + if [[ -f "${dest}" ]]; then + echo "[SKIP] ${fname}" + return 0 + fi + echo "[DOWNLOAD] ${fname}" + python3 - </dev/null | head -20 +fi diff --git a/examples/mem_agent/prepare_data.py b/examples/mem_agent/prepare_data.py new file mode 100644 index 000000000..595b63d7a --- /dev/null +++ b/examples/mem_agent/prepare_data.py @@ -0,0 +1,253 @@ +""" +Convert MemAgent hotpotqa parquet data to vime-compatible JSONL format. + +MemAgent parquet fields: + prompt : [{"role": "user", "content": question}] + context : "Document 1:\n...\n\nDocument 2:\n..." + reward_model : {"style": "rule", "ground_truth": ["answer1", ...]} + extra_info : {"index": 0, "question": ..., "num_docs": 200} + data_source : "hotpotqa" + ability : "memory" + +Output JSONL fields (vime convention): + prompt : question string (--input-key prompt) + label : first answer string (--label-key label) + metadata : { + "context" : long document text, + "ground_truth" : [all acceptable answers], # used by reward_func for multi-answer matching + "num_docs" : int, + "data_source" : str, + } + +Usage: + python prepare_data.py \\ + --input /path/to/hotpotqa_train.parquet \\ + --output /path/to/hotpotqa_train.jsonl + + # Can also pull directly from HuggingFace + python prepare_data.py \\ + --hf-dataset BytedTsinghua-SIA/hotpotqa \\ + --hf-split train \\ + --output /path/to/hotpotqa_train.jsonl +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +try: + import numpy as np +except ImportError: + np = None + + +def _to_json_safe(obj): + if np is not None: + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, (np.integer, np.floating, np.bool_)): + return obj.item() + if isinstance(obj, dict): + return {k: _to_json_safe(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_to_json_safe(v) for v in obj] + return obj + + +def convert_row(row: dict) -> dict | None: + """Convert one row from MemAgent format to vime JSONL format. + + Compatible with two formats: + Training set format (parquet): prompt(list) / reward_model / extra_info / context + Evaluation set format (eval_*.json): input / answers / num_docs / context + """ + context = row.get("context", "") + if not context: + return None + + # ── Evaluation set format: input + answers ─────────────────────────────── + if "input" in row: + question = row["input"] + answers = row.get("answers", []) + if isinstance(answers, str): + answers = [answers] + label = answers[0] if answers else "" + if not question or not label: + return None + return { + "prompt": question, + "label": label, + "metadata": { + "context": context, + "ground_truth": answers, + "num_docs": row.get("num_docs", 0), + "data_source": "hotpotqa", + }, + } + + # ── Training set format: prompt(list) + reward_model ──────────────────── + prompt_field = row.get("prompt", []) + if isinstance(prompt_field, list) and prompt_field: + question = prompt_field[0].get("content", "") if isinstance(prompt_field[0], dict) else str(prompt_field[0]) + elif isinstance(prompt_field, str): + question = prompt_field + else: + question = row.get("extra_info", {}).get("question", "") + + if not question: + return None + + reward_model = row.get("reward_model", {}) + if isinstance(reward_model, str): + try: + reward_model = json.loads(reward_model) + except Exception: + reward_model = {} + ground_truth = reward_model.get("ground_truth", []) + if isinstance(ground_truth, str): + ground_truth = [ground_truth] + label = ground_truth[0] if ground_truth else "" + if not label: + return None + + extra_info = row.get("extra_info", {}) or {} + + return { + "prompt": question, + "label": label, + "metadata": { + "context": context, + "ground_truth": ground_truth, + "num_docs": extra_info.get("num_docs", 0), + "data_source": row.get("data_source", "hotpotqa"), + }, + } + + +def convert_parquet(input_path: str, output_path: str) -> int: + try: + import pandas as pd + except ImportError: + print("ERROR: pandas is required. pip install pandas pyarrow", file=sys.stderr) + sys.exit(1) + + df = pd.read_parquet(input_path) + rows = df.to_dict(orient="records") + return _write_jsonl(rows, output_path) + + +def convert_hf(dataset_name: str, split: str, output_path: str) -> int: + try: + from datasets import load_dataset + except ImportError: + print("ERROR: datasets is required. pip install datasets", file=sys.stderr) + sys.exit(1) + + # Prefer the HF_ENDPOINT environment variable; otherwise automatically try the mirror + import os + + if not os.environ.get("HF_ENDPOINT"): + os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + + # Use list_repo_files to precisely list files for the target split + from huggingface_hub import list_repo_files + + exts = (".parquet", ".json", ".jsonl", ".csv") + + # list_repo_files is more reliable than glob and lists all files + all_files = list(list_repo_files(dataset_name, repo_type="dataset")) + split_files = [ + f"hf://datasets/{dataset_name}/{p}" for p in all_files if split in Path(p).name and Path(p).suffix in exts + ] + if split_files: + fmt = "parquet" if split_files[0].endswith(".parquet") else "json" + ds = load_dataset(fmt, data_files={split: split_files}, split=split) + else: + # Last fallback: the split name may be nested inside a directory (e.g. data/split/xxx.parquet) + split_files = [ + f"hf://datasets/{dataset_name}/{p}" + for p in all_files + if (f"/{split}/" in p or f"/{split}-" in p) and Path(p).suffix in exts + ] + if not split_files: + raise FileNotFoundError( + f"Cannot find files for split '{split}' in dataset '{dataset_name}'. " + f"Available files: {all_files[:20]}" + ) + fmt = "parquet" if split_files[0].endswith(".parquet") else "json" + ds = load_dataset(fmt, data_files={split: split_files}, split=split) + rows = [dict(r) for r in ds] + return _write_jsonl(rows, output_path) + + +def _write_jsonl(rows: list[dict], output_path: str) -> int: + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + written = 0 + skipped = 0 + with open(output_path, "w", encoding="utf-8") as f: + for row in rows: + out = convert_row(row) + if out is None: + skipped += 1 + continue + f.write(json.dumps(_to_json_safe(out), ensure_ascii=False) + "\n") + written += 1 + + print(f"Written: {written} Skipped: {skipped} → {output_path}") + return written + + +def convert_hf_file(dataset_name: str, filename: str, output_path: str) -> int: + """Directly download and convert a specified file from an HF repo; used for non-standard split files such as eval_*.json.""" + import os + + if not os.environ.get("HF_ENDPOINT"): + os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + + from huggingface_hub import hf_hub_download + + local_path = hf_hub_download( + repo_id=dataset_name, + filename=filename, + repo_type="dataset", + ) + suffix = Path(local_path).suffix.lower() + if suffix == ".parquet": + return convert_parquet(local_path, output_path) + else: + with open(local_path, encoding="utf-8") as f: + data = json.load(f) + if isinstance(data, dict): + rows = list(data.values()) if all(isinstance(v, dict) for v in data.values()) else [data] + else: + rows = data + return _write_jsonl(rows, output_path) + + +def main(): + parser = argparse.ArgumentParser(description="Convert MemAgent parquet to vime JSONL") + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--input", help="Local parquet file path") + group.add_argument("--hf-dataset", help="HuggingFace dataset name, e.g. BytedTsinghua-SIA/hotpotqa") + parser.add_argument("--hf-split", default="train", help="HF split (default: train)") + parser.add_argument( + "--hf-file", + default=None, + help="Directly specify a filename in the HF repo, e.g. eval_1600.json (for non-standard splits)", + ) + parser.add_argument("--output", required=True, help="Output JSONL file path") + args = parser.parse_args() + + if args.input: + convert_parquet(args.input, args.output) + elif args.hf_file: + convert_hf_file(args.hf_dataset, args.hf_file, args.output) + else: + convert_hf(args.hf_dataset, args.hf_split, args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/mem_agent/rollout.py b/examples/mem_agent/rollout.py new file mode 100644 index 000000000..4e13f9dbd --- /dev/null +++ b/examples/mem_agent/rollout.py @@ -0,0 +1,263 @@ +""" +MemAgent rollout for vime. + +Chunk-by-chunk memory update pipeline: + for chunk in split(context): + memory = LLM(problem, memory, chunk) + answer = LLM(problem, memory) # final turn with \\boxed{} +""" + +from __future__ import annotations + +import os +import traceback +from typing import Any + +from examples.mem_agent.rollout_client import MemAgentRolloutClient + +from vime.rollout.vllm_rollout import GenerateState +from vime.utils.types import Sample + +CHUNK_TOKENS = int(os.environ.get("MEM_CHUNK_TOKENS", "2048")) +MAX_MEMORY_TOKENS = int(os.environ.get("MEM_MAX_MEMORY", "1024")) +MAX_FINAL_TOKENS = int(os.environ.get("MEM_MAX_FINAL", "256")) +MAX_CHUNKS = int(os.environ.get("MEM_MAX_CHUNKS", "512")) + +_MEMORY_TEMPLATE = """You are presented with a problem, a section of an article that may contain the answer to the problem, and a previous memory. Please read the provided section carefully and update the memory with the new information that helps to answer the problem. Be sure to retain all relevant details from the previous memory while adding any new, useful information. + + +{prompt} + + + +{memory} + + +
+{chunk} +
+ +Updated memory: +""" + +_FINAL_TEMPLATE = """You are presented with a problem and a previous memory. Please answer the problem based on the previous memory and put the answer in \\boxed{{}}. + + +{prompt} + + + +{memory} + + +Your answer: +""" + +_NO_MEMORY = "No previous memory" +_STOP_TOKEN_STRINGS = ["", "<|endoftext|>"] + + +def _strip_stop_tokens(text: str) -> str: + for tok in _STOP_TOKEN_STRINGS: + text = text.replace(tok, "") + return text.strip() + + +async def generate( + args: Any, + sample: Sample, + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> Sample: + state = GenerateState(args) + client = MemAgentRolloutClient(args, state.tokenizer, sampling_params) + + if not isinstance(sample.metadata, dict): + sample.metadata = {} + + question = sample.prompt if isinstance(sample.prompt, str) else sample.prompt[-1]["content"] + context = sample.metadata.get("context", "") + if not context: + sample.status = Sample.Status.ABORTED + sample.rollout_log_probs = [] + return sample + + try: + context_ids = state.tokenizer.encode(context, add_special_tokens=False) + chunks_ids = [context_ids[i : i + CHUNK_TOKENS] for i in range(0, len(context_ids), CHUNK_TOKENS)][:MAX_CHUNKS] + + turns: list[dict] = [] + memory = _NO_MEMORY + mem_params = {**sampling_params, "max_new_tokens": MAX_MEMORY_TOKENS} + + for chunk_ids in chunks_ids: + chunk_text = state.tokenizer.decode(chunk_ids, skip_special_tokens=True) + messages = [ + { + "role": "user", + "content": _MEMORY_TEMPLATE.format(prompt=question, memory=memory, chunk=chunk_text), + } + ] + out = await client.generate(messages, sampling_params=mem_params) + memory = _strip_stop_tokens(out.response) or memory + turns.append( + { + "tokens": out.prompt_token_ids + out.token_ids, + "response_length": len(out.token_ids), + "loss_mask": [1] * len(out.token_ids), + "rollout_log_probs": out.log_probs, + } + ) + + final_messages = [ + { + "role": "user", + "content": _FINAL_TEMPLATE.format(prompt=question, memory=memory), + } + ] + final_out = await client.generate( + final_messages, + sampling_params={**sampling_params, "max_new_tokens": MAX_FINAL_TOKENS}, + ) + turns.append( + { + "tokens": final_out.prompt_token_ids + final_out.token_ids, + "response_length": len(final_out.token_ids), + "loss_mask": [1] * len(final_out.token_ids), + "rollout_log_probs": final_out.log_probs, + } + ) + + sample.metadata["final_output"] = _strip_stop_tokens(final_out.response) + sample.train_metadata = {"turns": turns} + + if os.environ.get("MEM_DEBUG"): + print( + f"[MemAgent] {len(chunks_ids)} chunks -> {len(turns)} turns, " + f"response_lengths={[t['response_length'] for t in turns]}" + ) + + first = turns[0] + first_prompt_len = len(first["tokens"]) - first["response_length"] + prompt_token_ids = first["tokens"][:first_prompt_len] + + cat_token_ids: list[int] = [] + cat_loss_mask: list[int] = [] + cat_log_probs: list[float] = [] + for t in turns: + p_len = len(t["tokens"]) - t["response_length"] + cat_token_ids += t["tokens"] + cat_loss_mask += [0] * p_len + t["loss_mask"] + cat_log_probs += [0.0] * p_len + t["rollout_log_probs"] + + sample.prompt = final_messages[0]["content"] + sample.response = _strip_stop_tokens(final_out.response) + sample.tokens = prompt_token_ids + cat_token_ids + sample.response_length = len(cat_token_ids) + sample.loss_mask = cat_loss_mask + sample.rollout_log_probs = cat_log_probs + sample.status = Sample.Status.TRUNCATED if final_out.finish_reason == "length" else Sample.Status.COMPLETED + + except Exception: + traceback.print_exc() + sample.response = "" + sample.rollout_log_probs = [] + sample.status = Sample.Status.FAILED + + return sample + + +def _last_boxed_only_string(string: str) -> str | None: + if "\\boxed " in string: + return "\\boxed " + string.split("\\boxed ")[-1].split("$")[0] + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + i, right_brace_idx, num_left_braces_open = idx, None, 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + return string[idx : right_brace_idx + 1] if right_brace_idx is not None else None + + +def _remove_boxed(s: str) -> str: + if s.startswith("\\boxed "): + return s[len("\\boxed ") :] + left = "\\boxed{" + assert s.startswith(left) and s.endswith("}") + return s[len(left) : -1] + + +def _extract_boxed(text: str) -> str: + s = _last_boxed_only_string(text) + if s is None: + return "" + try: + return _remove_boxed(s).strip() + except Exception: + return "" + + +def _strip_string(string: str) -> str: + string = string.replace("\n", "") + string = string.replace("\\!", "") + string = string.replace("\\\\", "\\") + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + string = string.replace("\\left", "") + string = string.replace("\\right", "") + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + string = string.replace("\\$", "") + string = string.replace("\\%", "") + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + if len(string) == 0: + return string + return string.replace(" ", "") + + +def _is_equiv(str1: str, str2: str) -> bool: + if str1 is None and str2 is None: + return True + if str1 is None or str2 is None: + return False + try: + return _strip_string(str1) == _strip_string(str2) + except Exception: + return str1 == str2 + + +async def reward_func(args: Any, sample: Any, **kwargs) -> dict: + metadata = sample.metadata if isinstance(sample.metadata, dict) else {} + final_output = metadata.get("final_output", "") or sample.response or "" + ground_truth = metadata.get("ground_truth", []) + if not ground_truth: + label = str(sample.label) if sample.label is not None else "" + ground_truth = [label] if label else [] + + solution_str = final_output[-300:] + pred = _extract_boxed(solution_str) or "" + + score = 0.0 + for gt in ground_truth: + gt_lower = gt.lower() + try: + boxed = _last_boxed_only_string(solution_str) + if boxed is not None: + answer = _remove_boxed(boxed) + if _is_equiv(answer.lower(), gt_lower): + score = 1.0 + break + except Exception: + pass + + return {"score": score, "pred": pred, "gt": ground_truth[0] if ground_truth else ""} diff --git a/examples/mem_agent/rollout_client.py b/examples/mem_agent/rollout_client.py new file mode 100644 index 000000000..e492a9bbf --- /dev/null +++ b/examples/mem_agent/rollout_client.py @@ -0,0 +1,79 @@ +"""Lightweight vLLM router client for MemAgent multi-turn rollouts.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from vime.rollout.vllm_rollout import _build_inference_sampling_params, _inference_generate_tokens_and_logprobs +from vime.utils.http_utils import post + + +@dataclass +class GenerationOutput: + prompt_text: str + prompt_token_ids: list[int] + response: str + token_ids: list[int] + log_probs: list[float] + finish_reason: str + + +class MemAgentRolloutClient: + """Wrap vLLM ``/inference/v1/generate`` for chat-style MemAgent turns.""" + + def __init__(self, args: Any, tokenizer: Any, sampling_params: dict): + self.base = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + self.model = args.hf_checkpoint + self.tokenizer = tokenizer + self.sampling_params = dict(sampling_params) + + async def generate( + self, + messages: list[dict[str, str]], + sampling_params: dict | None = None, + ) -> GenerationOutput: + params = dict(self.sampling_params) + if sampling_params: + params.update(sampling_params) + + prompt_text = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + prompt_token_ids = self.tokenizer(prompt_text, add_special_tokens=False)["input_ids"] + + payload = { + "model": self.model, + "token_ids": prompt_token_ids, + "sampling_params": _build_inference_sampling_params(params), + } + output = await post(f"{self.base}/inference/v1/generate", payload) + choice = output["choices"][0] + + skip_sp = params.get("skip_special_tokens") + skip_decode = True if skip_sp is None else bool(skip_sp) + out_ids = choice.get("token_ids") or [] + text = ( + self.tokenizer.decode(out_ids, skip_special_tokens=skip_decode) + if isinstance(out_ids, list) and out_ids + else "" + ) + + token_ids, log_probs = _inference_generate_tokens_and_logprobs(choice) + + fr = choice.get("finish_reason") or "stop" + if isinstance(fr, dict): + finish_reason = fr.get("type", "stop") + else: + finish_reason = "length" if fr == "length" else "stop" + + return GenerationOutput( + prompt_text=prompt_text, + prompt_token_ids=prompt_token_ids, + response=text, + token_ids=token_ids, + log_probs=log_probs, + finish_reason=finish_reason, + ) diff --git a/examples/mem_agent/run-eval.sh b/examples/mem_agent/run-eval.sh new file mode 100644 index 000000000..1aa52eab5 --- /dev/null +++ b/examples/mem_agent/run-eval.sh @@ -0,0 +1,161 @@ +#!/usr/bin/env bash +# MemAgent RULER-HQA evaluation via vLLM serve + eval_ruler_hqa.py. +# +# Usage: +# MODEL_PATH=/path/to/hf_ckpt bash examples/mem_agent/run-eval.sh +# CONVERT=1 SINGLE_ITER=iter_0000199 bash examples/mem_agent/run-eval.sh +# MODEL_PATH=/root/Qwen3-4B SAVE_FILE=Qwen3-4B-base bash examples/mem_agent/run-eval.sh +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# shellcheck source=_common.sh +source "${SCRIPT_DIR}/_common.sh" + +EVAL_PY="${MEM_AGENT_DIR}/eval_ruler_hqa.py" +TP="${TP:-1}" +SERVE_HOST="${SERVE_HOST:-127.0.0.1}" +SERVE_PORT="${SERVE_PORT:-8000}" +LENGTH="${LENGTH:-50 200 800}" +SAVE_DIR="${SAVE_DIR:-${MEM_AGENT_DIR}/results}" +API="${API:-recurrent}" +N_PROC="${N_PROC:-16}" +TEMPERATURE="${TEMPERATURE:-0.7}" +TOP_P="${TOP_P:-0.95}" +FORCE="${FORCE:-0}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}" +GPU_MEMORY_UTIL="${GPU_MEMORY_UTIL:-0.85}" + +CHECKPOINT_DIR="${CHECKPOINT_DIR:-${SAVE_PATH}}" +SINGLE_ITER="${SINGLE_ITER:-iter_0000199}" +CONVERT="${CONVERT:-0}" + +STAMP=$(date +%Y%m%d_%H%M%S) +LOG_FILE="${SAVE_DIR}/eval_${STAMP}.log" +mkdir -p "${SAVE_DIR}" + +exec > >(tee -a "${LOG_FILE}") 2>&1 +echo "=== MemAgent eval started at $(date) ===" +echo "LOG_FILE=${LOG_FILE}" + +if [[ "${CONVERT}" == "1" ]]; then + echo "[step] Converting ${SINGLE_ITER} ..." + SINGLE_ITER="${SINGLE_ITER}" CHECKPOINT_DIR="${CHECKPOINT_DIR}" \ + bash "${MEM_AGENT_DIR}/convert-to-hf.sh" + MODEL_PATH="${MODEL_PATH:-${CHECKPOINT_DIR}-HF/${SINGLE_ITER}}" +fi + +if [[ -z "${MODEL_PATH:-}" ]]; then + echo "ERROR: MODEL_PATH is required." + echo " MODEL_PATH=/path/to/hf_ckpt bash examples/mem_agent/run-eval.sh" + echo " CONVERT=1 bash examples/mem_agent/run-eval.sh" + exit 1 +fi + +if [[ ! -d "${MODEL_PATH}" ]]; then + echo "ERROR: MODEL_PATH not found: ${MODEL_PATH}" + exit 1 +fi + +SAVE_FILE="${SAVE_FILE:-$(basename "${MODEL_PATH%/}")}" +MODEL_NAME="${MODEL_PATH}" + +LENGTHS="${LENGTH}" bash "${MEM_AGENT_DIR}/prepare-eval-data.sh" + +export PYTHONPATH="/root/Megatron-LM:${VIME_ROOT}:${PYTHONPATH:-}" +export VLLM_SERVE_HOST="${SERVE_HOST}" VLLM_SERVE_PORT="${SERVE_PORT}" +export SERVE_HOST SERVE_PORT +export DATAROOT="${DATA_ROOT}" + +log() { echo "[$(date '+%H:%M:%S')] $*"; } + +wait_for_server() { + local url="http://${SERVE_HOST}:${SERVE_PORT}/v1/models" + log "Waiting for vLLM at ${url} ..." + local attempts=0 + local max_attempts=120 + while true; do + if [[ -n "${VLLM_PID:-}" ]] && ! kill -0 "${VLLM_PID}" 2>/dev/null; then + log "ERROR: vLLM process (PID ${VLLM_PID}) died. See ${VLLM_LOG:-server log}." + exit 1 + fi + resp=$(curl -sf --max-time 10 "${url}" 2>/dev/null || true) + if echo "${resp}" | grep -Fq "${MODEL_NAME}" 2>/dev/null; then + log "vLLM ready." + break + fi + attempts=$((attempts + 1)) + if (( attempts >= max_attempts )); then + log "ERROR: vLLM not ready after ${max_attempts} attempts." + exit 1 + fi + if (( attempts % 6 == 0 )); then + found=$(echo "${resp}" | grep -o '"id":"[^"]*"' 2>/dev/null | head -3 || echo "(no response)") + log "Still waiting... models: ${found}" + fi + sleep 5 + done +} + +kill_server() { + if [[ -n "${VLLM_PID:-}" ]]; then + log "Stopping vLLM (pid=${VLLM_PID}) ..." + kill -TERM "${VLLM_PID}" 2>/dev/null || true + wait "${VLLM_PID}" 2>/dev/null || true + fi +} + +build_common_args() { + local extra=() + extra+=(--model "${MODEL_NAME}") + extra+=(--tokenizer "${MODEL_PATH}") + extra+=(--api "${API}") + extra+=(--n-proc "${N_PROC}") + extra+=(--temperature "${TEMPERATURE}") + extra+=(--top-p "${TOP_P}") + if [[ "${FORCE}" == "1" ]]; then + extra+=(--force) + fi + echo "${extra[@]}" +} + +run_hqa() { + local common + read -ra common <<< "$(build_common_args)" + for length in ${LENGTH}; do + local subdir="${SAVE_DIR}/ruler_hqa_${length}" + log "==> ruler_hqa n_docs=${length}" + python3 "${EVAL_PY}" \ + "${common[@]}" \ + --length "${length}" \ + --data-root "${DATA_ROOT}" \ + --save-dir "${subdir}" \ + --save-file "${SAVE_FILE}" + done +} + +log "MODEL_PATH=${MODEL_PATH}" +log "TP=${TP} LENGTH=${LENGTH} N_PROC=${N_PROC}" +log "DATA_ROOT=${DATA_ROOT} SAVE_DIR=${SAVE_DIR}" + +pkill -9 -f '[v]llm serve' 2>/dev/null || true +sleep 2 + +VLLM_LOG="${SAVE_DIR}/vllm_server_${STAMP}.log" +log "Starting vLLM serve (tp=${TP}, port=${SERVE_PORT}, max_len=${MAX_MODEL_LEN}) ..." + +CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" \ + vllm serve "${MODEL_PATH}" \ + --tensor-parallel-size "${TP}" \ + --host "${SERVE_HOST}" \ + --port "${SERVE_PORT}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --gpu-memory-utilization "${GPU_MEMORY_UTIL}" \ + --trust-remote-code \ + > "${VLLM_LOG}" 2>&1 & +VLLM_PID=$! +trap 'kill_server' EXIT INT TERM + +wait_for_server +run_hqa + +log "=== Eval finished. Results: ${SAVE_DIR} ===" diff --git a/examples/mem_agent/run-qwen3-4b-train.sh b/examples/mem_agent/run-qwen3-4b-train.sh new file mode 100644 index 000000000..119038977 --- /dev/null +++ b/examples/mem_agent/run-qwen3-4b-train.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# MemAgent GRPO training — Qwen3-4B, vime + vLLM colocate. +# +# Usage: +# bash examples/mem_agent/run-qwen3-4b-train.sh +# NUM_ROLLOUT=200 SAVE_PATH=/path/to/save bash examples/mem_agent/run-qwen3-4b-train.sh +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# shellcheck source=_common.sh +source "${SCRIPT_DIR}/_common.sh" + +export NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +export ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +export N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-64}" +export SAVE_INTERVAL="${SAVE_INTERVAL:-50}" + +echo "=== MemAgent train NUM_ROLLOUT=${NUM_ROLLOUT} SAVE_PATH=${SAVE_PATH} ===" +if [[ ! -f "${TRAIN_DATA}" ]]; then + echo "ERROR: training data not found: ${TRAIN_DATA}" + exit 1 +fi + +mem_agent_detect_nvlink +mem_agent_detect_gpus +mem_agent_cleanup +mem_agent_launch_train + +echo "=== MemAgent train finished ===" diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py index ec026860e..c04640e6f 100644 --- a/examples/multi_agent/agent_system.py +++ b/examples/multi_agent/agent_system.py @@ -5,28 +5,13 @@ from copy import deepcopy from vime.rollout.rm_hub import batched_async_rm -from vime.rollout.vllm_rollout import _build_inference_sampling_params +from vime.rollout.vllm_rollout import _build_inference_sampling_params, _inference_generate_tokens_and_logprobs from vime.utils.http_utils import post from vime.utils.types import Sample from .prompts import SOLVER_PROMPT_TEMPLATE, generate_rewriter_template, generate_select_template -def _inference_generate_tokens_and_logprobs(choice): - """Compatibility helper for current VIME vLLM rollout responses.""" - new_response_tokens = choice.get("token_ids") or [] - new_response_log_probs = [] - lp = choice.get("logprobs") - if isinstance(lp, dict): - content_items = lp.get("content") or [] - new_response_log_probs = [ - float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items - ] - if not new_response_log_probs: - new_response_log_probs = [0.0] * len(new_response_tokens) - return new_response_tokens, new_response_log_probs - - async def generate_response(args, prompt, key, worker_id: int | None = None): try: sampling_params = args.sampling_params @@ -66,12 +51,12 @@ async def generate_response(args, prompt, key, worker_id: int | None = None): new_response_tokens, new_response_log_probs = _inference_generate_tokens_and_logprobs(choice) response_text = tokenizer.decode(new_response_tokens, skip_special_tokens=False) if new_response_tokens else "" - # Update sample with tokens directly - avoiding re-tokenization - sample.tokens = sample.tokens + new_response_tokens - sample.response_length += len(new_response_tokens) - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += new_response_log_probs + sample.append_response_tokens( + args, + tokens=new_response_tokens, + log_probs=new_response_log_probs, + trainable=True, + ) assert len(sample.rollout_log_probs) == sample.response_length, ( f"rollout logprob length mismatch: {len(sample.rollout_log_probs)} logprobs " f"vs {sample.response_length} response tokens" diff --git a/examples/multi_agent/run-qwen3-30B-A3B-multi-agent.sh b/examples/multi_agent/run-qwen3-30B-A3B-multi-agent.sh index a1d2603fc..efb6b6ef2 100644 --- a/examples/multi_agent/run-qwen3-30B-A3B-multi-agent.sh +++ b/examples/multi_agent/run-qwen3-30B-A3B-multi-agent.sh @@ -1,7 +1,7 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray diff --git a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh index 9b8424be3..b03217f55 100755 --- a/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh +++ b/examples/multi_agent/run-qwen3-4B-multi-agent-npu.sh @@ -44,7 +44,6 @@ source "${SCRIPT_DIR}/models/qwen3-4B.sh" CKPT_ARGS=( --hf-checkpoint "${WEIGHT_DIR}" --load "${WEIGHT_DIR}" - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -56,8 +55,6 @@ ROLLOUT_ARGS=( --rollout-shuffle --rm-type math - --rollout-backend vllm - --vllm-weight-sync-mode native --vllm-gpu-memory-utilization 0.6 --vllm-enable-sleep-mode --vllm-max-model-len 4096 @@ -116,6 +113,7 @@ WANDB_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 ) @@ -126,11 +124,10 @@ MISC_ARGS=( --attention-softmax-in-fp32 --attention-backend flash --use-flash-attn - --train-memory-margin-bytes 2147483648 ) export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --num-gpus 0 --resources '{"NPU": 8}' --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 ray job submit --address="http://127.0.0.1:8265" \ -- python3 train.py \ diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md new file mode 100644 index 000000000..9d85d3e69 --- /dev/null +++ b/examples/on_policy_distillation/README.md @@ -0,0 +1,168 @@ +# On-Policy Distillation Example + +This example shows how to run **on-policy distillation (OPD)** using vime. A +small student (Qwen3-8B) is aligned to imitate a larger teacher (Qwen3-32B) by +training only on the student's own rollouts and matching the teacher's +token-level log-probabilities. + +## Key Features + +- **OPD is orthogonal to advantage estimators**: OPD works as an additive KL + penalty on top of any advantage estimator (GRPO, PPO, REINFORCE++, etc.), not + as a separate estimator. +- **Two teacher modes**: + - **vllm**: Teacher runs on an external vLLM server; teacher log-probs are + obtained during rollout. + - **megatron**: Teacher is loaded directly into Megatron via + `--opd-teacher-load`; teacher log-probs are computed during the training + forward pass. +- **Student rollout always uses vLLM** (vime's default rollout backend). + +## Key Arguments + +| Argument | Description | +|----------|-------------| +| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. | +| `--opd-type` | Type of OPD: `vllm` or `megatron`. Required when `--use-opd` is set. | +| `--opd-kl-coef` | OPD KL penalty coefficient (default: 1.0). | +| `--opd-teacher-load` | Path to teacher checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=vllm`. | +| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. | + +## Mode Comparison + +| Mode | Teacher Location | When to use | +|------|------------------|-------------| +| `vllm` | External vLLM server | Teacher has different architecture or is larger than GPU memory | +| `megatron` | Loaded into Megatron training | Teacher has same architecture as policy/ref model | + +## Components + +- `vime/rollout/on_policy_distillation.py` implements (for vLLM mode): + - `reward_func` calls the teacher server (via `args.rm_url`) with every sample + to obtain token-level logprobs. + - `post_process_rewards` trims the teacher logprobs to the generated response + span and writes the tensors back to each `Sample` to compute advantages. +- `run-qwen3-8B-opd.sh` launches a vLLM teacher server, then submits a Ray job + that runs `train.py`. +- `run-qwen3-8B-opd-megatron.sh` uses a Megatron-loaded teacher model (no + external server needed). + +## Running the example + +### Using vLLM Teacher (External Server) + +1. Download or prepare the required checkpoints and data. + +```bash +hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B +hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k +``` + +2. Run the hf to mcore for student model conversion: + +```bash +cd /root/vime +source scripts/models/qwen3-8B.sh + +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen3-8B \ + --save /root/Qwen3-8B_torch_dist +``` + +3. Run on-policy distillation: + +```bash +bash examples/on_policy_distillation/run-qwen3-8B-opd.sh +``` + +GPU layout: + +| GPUs | Role | +|------|------| +| 0–3 | Student Megatron train + student vLLM rollout (colocate) | +| 4–7 | Teacher vLLM (Qwen3-32B, TP=4) | + +### Using Megatron Teacher (No External Server) + +1. Prepare student checkpoint (same as above). + +2. **IMPORTANT**: Convert your teacher model to Megatron format (change the path + to your actual teacher): + +```bash +# This example uses the same model as both student and teacher (for demonstration only) +# In practice, use a different (stronger) model as the teacher! +cd /root/vime +source scripts/models/qwen3-8B.sh # Or your teacher model config + +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/YourTeacherModel \ + --save /root/YourTeacherModel_torch_dist +``` + +3. Edit `run-qwen3-8B-opd-megatron.sh` to update paths: + - Change `--opd-teacher-load` to your teacher model path + - Adjust `--opd-kl-coef` based on your task + +4. Run: + +```bash +bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh +``` + +# Preliminary Results + +End-to-end run with `run-qwen3-8B-opd.sh` (dapo-math-17k train, GRPO + +`--opd-kl-coef 1.0`, ~220 rollouts / iter_0000219). Offline GSM8K greedy eval: + +| Model | GSM8K Accuracy | +|-------|----------------| +| Qwen3-8B (pre-OPD) | 79.7% (n=300) | +| Qwen3-8B (post-OPD) | **88.2%** (n=1319, **+8.5 pp**) | +| Qwen3-32B teacher | 87.0% (n=300) | + +Training health signal: `rollout/opd_reverse_kl` dropped from 0.145 → ~0.10 +(−38%). Pure OPD uses `raw_reward=0`; the learning signal is the OPD KL term. + +# FAQ + +1. **Why are there two OPD modes?** + - `vllm` mode: The teacher runs on an independent vLLM server. This is useful + when the teacher has a different architecture or is too large to load + together with the policy model. + - `megatron` mode: The teacher is loaded into Megatron using the same + parameter loading mechanism as the reference model. This requires the + teacher to have the same architecture as the policy model. + +2. **How do I use Megatron-based teacher instead of vLLM server?** + Replace your OPD arguments: + ```bash + # Instead of: + --use-opd --opd-type vllm --opd-kl-coef 1.0 + # Use: + --use-opd --opd-type megatron --opd-kl-coef 1.0 --opd-teacher-load /path/to/teacher_checkpoint + ``` + +3. **What happens if I set wrong arguments?** + The system will raise clear errors: + - `--use-opd` without `--opd-type`: Error asking you to specify type + - `--opd-type megatron` without `--opd-teacher-load`: Error asking for teacher checkpoint + - `--opd-type vllm` with `--opd-teacher-load`: Error indicating conflict + +4. **Why is `rollout/raw_reward` always 0?** + Pure OPD distillation does not use an external reward model. The learning + signal comes entirely from the OPD KL term applied to advantages. + +5. **Self-distillation: why is `opd_reverse_kl` near 0 at the start?** + Teacher and student start from the same weights, so reverse KL is ~0 until + the student updates. For a true distillation signal, use a stronger / + differently trained teacher (or `--opd-type vllm` with Qwen3-32B). + +# References + +1. https://thinkingmachines.ai/blog/on-policy-distillation/ +2. https://arxiv.org/abs/2306.13649 +3. https://arxiv.org/abs/2306.08543 diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh new file mode 100644 index 000000000..40c3ae64e --- /dev/null +++ b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh @@ -0,0 +1,161 @@ +#!/bin/bash + +# On-Policy Distillation with Megatron-based teacher model +# This example uses the original model as the teacher (self-distillation for demonstration) +# +# IMPORTANT: This is just an example configuration! +# In practice, you should: +# 1. Use a different (stronger) model as the teacher +# 2. Adjust --opd-kl-coef based on your task +# 3. Configure proper evaluation metrics + +set -ex + +export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +source "/root/vime/scripts/models/qwen3-8B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-8B + --ref-load /root/Qwen3-8B_torch_dist + --load /root/Qwen3-8B_torch_dist + --save /root/Qwen3-8B_vime/ + --save-interval 10 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 300 + --rollout-batch-size 16 + --n-samples-per-prompt 4 + --rollout-max-response-len 16384 + --rollout-temperature 1 + + --global-batch-size 64 + --balance-data +) + +RM_ARGS=( + --rm-type math +) + +EVAL_ARGS=( + # --eval-interval 20 + # --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + # --n-samples-per-eval-prompt 16 + # --eval-max-response-len 16384 + # --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 36 + + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator grpo + # OPD Configuration + --use-opd + --opd-type megatron + --opd-kl-coef 1.0 + # CHANGE THIS to a stronger teacher checkpoint in practice + --opd-teacher-load /root/Qwen3-8B_torch_dist + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-opd + # --wandb-group qwen3-8B-opd-megatron + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.4 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --make-vocab-size-divisible-by 128 +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --working-dir /root/vime \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${RM_ARGS[@]} + +#### clear after training +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 -f "train.py" || true +sleep 3 diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd.sh b/examples/on_policy_distillation/run-qwen3-8B-opd.sh new file mode 100644 index 000000000..4c86f3473 --- /dev/null +++ b/examples/on_policy_distillation/run-qwen3-8B-opd.sh @@ -0,0 +1,201 @@ +#!/bin/bash + +# usage: bash examples/on_policy_distillation/run-qwen3-8B-opd.sh + +set -ex + +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +# Start the teacher model server +TEACHER_IP="127.0.0.1" +TEACHER_PORT=13141 +LOG_FILE="/tmp/vllm_teacher_$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6).log" + +## Launch the teacher model server in the background +CUDA_VISIBLE_DEVICES=4,5,6,7 python3 -m vllm.entrypoints.openai.api_server \ + --model /root/Qwen3-32B \ + --host 0.0.0.0 \ + --port ${TEACHER_PORT} \ + --tensor-parallel-size 4 \ + --gpu-memory-utilization 0.85 \ + --trust-remote-code \ + --dtype bfloat16 \ + --max-model-len 16384 \ + --disable-custom-all-reduce \ + > "${LOG_FILE}" 2>&1 & +TEACHER_PID=$! + +echo "Starting teacher model server (pid=${TEACHER_PID})..." + +## Wait for the teacher model server to be ready +for i in $(seq 1 120); do + if ! kill -0 "${TEACHER_PID}" 2>/dev/null; then + echo "ERROR: Teacher server process died. Check ${LOG_FILE}" + tail -n 20 "${LOG_FILE}" + exit 1 + fi + HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://${TEACHER_IP}:${TEACHER_PORT}/health" 2>/dev/null || true) + if [ "${HTTP_CODE}" = "200" ]; then + echo "Teacher model server is up and running at ${TEACHER_IP}:${TEACHER_PORT}." + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Teacher server failed to start within 10 minutes" + tail -n 20 "${LOG_FILE}" + kill "${TEACHER_PID}" 2>/dev/null || true + exit 1 + fi + echo "Waiting for the teacher model server to start..." + sleep 5 +done +sleep 5 + +source "/root/vime/scripts/models/qwen3-8B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-8B + --ref-load /root/Qwen3-8B_torch_dist + --load /root/Qwen3-8B_torch_dist + --save /root/Qwen3-8B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 300 + --rollout-batch-size 16 + --n-samples-per-prompt 4 + --rollout-max-response-len 4096 + --rollout-max-context-len 8192 + --rollout-temperature 1 + + --global-batch-size 64 + --balance-data +) + +RM_ARGS=( + --custom-rm-path vime.rollout.on_policy_distillation.reward_func + --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards + --rm-url http://${TEACHER_IP}:${TEACHER_PORT}/inference/v1/generate +) + +EVAL_ARGS=( + # --eval-interval 50 + # --eval-prompt-data gsm8k /root/gsm8k/test.parquet + # --eval-input-key messages + # --n-samples-per-eval-prompt 1 + # --eval-max-response-len 4096 + # --eval-top-k 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 2048 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-opd + --opd-type vllm + --opd-kl-coef 1.0 + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-opd + # --wandb-group qwen3-8B-opd + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.4 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --make-vocab-size-divisible-by 128 +) + +# launch the master node of ray in container +export CUDA_VISIBLE_DEVICES=0,1,2,3 +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --working-dir /root/vime \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${RM_ARGS[@]} + +#### clear after training +kill ${TEACHER_PID} 2>/dev/null || true +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 -f "train.py" || true +sleep 3 diff --git a/examples/retool/retool_qwen3_4b_rl.sh b/examples/retool/retool_qwen3_4b_rl.sh index 42d4eb2da..9f3201b51 100644 --- a/examples/retool/retool_qwen3_4b_rl.sh +++ b/examples/retool/retool_qwen3_4b_rl.sh @@ -65,7 +65,6 @@ CKPT_ARGS=( --save /path/to/Qwen3-4B_vime_npu/ --save-interval 20 --no-load-optim - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -131,10 +130,10 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 --vllm-gpu-memory-utilization 0.7 --vllm-enable-sleep-mode - --vllm-weight-sync-mode native --vllm-max-model-len 16384 ) @@ -166,7 +165,7 @@ ray start --head \ --dashboard-host=0.0.0.0 # Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON=$(cat << 'EOF' +RUNTIME_ENV_JSON=$(cat << EOF { "env_vars": { "PYTHONPATH": "${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge", diff --git a/examples/retool/retool_qwen3_4b_sft.sh b/examples/retool/retool_qwen3_4b_sft.sh index cf067ab80..c556052ee 100644 --- a/examples/retool/retool_qwen3_4b_sft.sh +++ b/examples/retool/retool_qwen3_4b_sft.sh @@ -65,7 +65,6 @@ CKPT_ARGS=( --save-interval 1000 --save-hf /path/to/Qwen3-4B_sft_vime_hf/ --no-load-optim - --megatron-to-hf-mode bridge ) SFT_ARGS=( @@ -135,7 +134,7 @@ ray start --head \ --dashboard-host=0.0.0.0 # Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON=$(cat << 'EOF' +RUNTIME_ENV_JSON=$(cat << EOF { "env_vars": { "PYTHONPATH": "${VIME_DIR}:${VIME_DIR}/examples/retool:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge", diff --git a/examples/search-r1/run_qwen3_4b_npu.sh b/examples/search-r1/run_qwen3_4b_npu.sh index 8f2aa6dea..ed3b279da 100644 --- a/examples/search-r1/run_qwen3_4b_npu.sh +++ b/examples/search-r1/run_qwen3_4b_npu.sh @@ -66,7 +66,6 @@ CKPT_ARGS=( --save /path/to/Qwen3-4B-Instruct-2507_vime_npu/ --save-interval 100 --no-load-optim - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -135,10 +134,10 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 --vllm-gpu-memory-utilization 0.7 --vllm-enable-sleep-mode - --vllm-weight-sync-mode native ) MISC_ARGS=( @@ -173,7 +172,7 @@ ray start --head \ --dashboard-host=0.0.0.0 # Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON=$(cat << 'EOF' +RUNTIME_ENV_JSON=$(cat << EOF { "env_vars": { "PYTHONPATH": "${VIME_DIR}/examples/search-r1:/root/Megatron-LM:/root/vllm:/root/vllm-ascend:${VIME_DIR}:/root/Megatron-Bridge:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:/usr/local/Ascend/ascend-toolkit/latest/tools/ms_fmk_transplt/torch_npu_bridge", diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh new file mode 100644 index 000000000..1603c4f3f --- /dev/null +++ b/examples/tau-bench/run_qwen3_4B.sh @@ -0,0 +1,151 @@ +#!/bin/bash + +if grep -q $'\r' "$0" 2>/dev/null; then + exec bash <(sed 's/\r$//' "$0") "$@" +fi + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' 2>/dev/null || true +sleep 3 +ray stop --force 2>/dev/null || true +pkill -9 ray 2>/dev/null || true +pkill -9 -f 'python3 train.py' 2>/dev/null || true +sleep 3 + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B-Instruct-2507/ + --ref-load /root/Qwen3-4B-Instruct-2507_torch_dist/ + --save /root/Qwen3-4B-Instruct-2507_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/tau-bench/retail_train_tasks.jsonl + --input-key index + --rollout-shuffle + --num-rollout 500 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-max-context-len 16384 + --rollout-temperature 0.7 + --global-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + --balance-data +) + +EVAL_ARGS=( + --eval-interval 5 + --eval-prompt-data retail-dev /root/tau-bench/retail_dev_tasks.jsonl + --n-samples-per-eval-prompt 1 + --eval-max-response-len 4096 + --eval-top-k 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.001 + --kl-loss-type low_var_kl + --entropy-coef 0.01 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 5e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.7 + --vllm-max-model-len 16384 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +CUSTOM_ARGS=( + --custom-generate-function-path generate_with_tau.generate + --custom-rm-path generate_with_tau.batched_tau_bench_rm +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +NUM_GPUS=2 + +ray start --head \ + --node-ip-address "${MASTER_ADDR}" \ + --num-gpus "${NUM_GPUS}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC\": \"900\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node "${NUM_GPUS}" \ + --rollout-num-gpus "${NUM_GPUS}" \ + --colocate \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${CUSTOM_ARGS[@]}" \ + "${MISC_ARGS[@]}" diff --git a/examples/tau-bench/run_qwen3_4B_npu.sh b/examples/tau-bench/run_qwen3_4B_npu.sh index c5bad23fc..2c97f2eeb 100644 --- a/examples/tau-bench/run_qwen3_4B_npu.sh +++ b/examples/tau-bench/run_qwen3_4B_npu.sh @@ -41,7 +41,6 @@ CKPT_ARGS=( --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/ --load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/ --ref-load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/ - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -101,6 +100,7 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 1 --vllm-gpu-memory-utilization 0.7 --vllm-max-model-len 16384 @@ -123,7 +123,7 @@ CUSTOM_ARGS=( export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head \ +ray start --head --num-gpus 0 --resources '{"NPU": 8}' \ --node-ip-address "${MASTER_ADDR}" \ --disable-usage-stats \ --dashboard-host=0.0.0.0 \ @@ -145,4 +145,3 @@ ray job submit --address="http://127.0.0.1:8265" \ "${VLLM_ARGS[@]}" \ "${CUSTOM_ARGS[@]}" \ "${MISC_ARGS[@]}" - diff --git a/examples/tau-bench/token_delta.py b/examples/tau-bench/token_delta.py new file mode 100644 index 000000000..b78023d8c --- /dev/null +++ b/examples/tau-bench/token_delta.py @@ -0,0 +1,58 @@ +from typing import Any + + +def get_token_delta( + tokenizer: Any, + messages: list[dict[str, Any]], + *, + include_generation_prompt: bool = False, +) -> tuple[list[int], list[int]]: + """Return the tokens and loss mask contributed by the last chat message.""" + if not messages: + raise ValueError("Cannot calculate a token delta for an empty conversation") + + is_assistant = messages[-1]["role"] == "assistant" + curr = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False) + + if is_assistant: + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) + generation_prompt = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False) + if not generation_prompt.startswith(prev): + raise ValueError("Adding the assistant generation prompt rewrote the rendered conversation") + if not curr.startswith(generation_prompt): + raise ValueError("The assistant response does not extend its generation prompt") + + generation_prompt_text = generation_prompt[len(prev) :] + if not include_generation_prompt: + new_text = curr[len(generation_prompt) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + return new_tokens, [1] * len(new_tokens) + + new_text = curr[len(prev) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + generation_prompt_length = len(tokenizer.encode(generation_prompt_text, add_special_tokens=False)) + masked_prefix_length = min(generation_prompt_length, len(new_tokens)) + loss_mask = [0] * masked_prefix_length + loss_mask.extend([1] * (len(new_tokens) - masked_prefix_length)) + return new_tokens, loss_mask + + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) + + if curr.startswith(prev): + new_text = curr[len(prev) :] + elif messages[-1]["role"] == "user": + # Reasoning templates such as Qwen3 can rewrite history when a new user + # message arrives. Render that message independently instead of slicing + # the rewritten conversation at the old conversation length. + new_text = tokenizer.apply_chat_template( + [messages[-1]], + add_generation_prompt=False, + tokenize=False, + ) + if not curr.endswith(new_text): + raise ValueError("The latest user message is not a standalone suffix of the rendered conversation") + else: + raise ValueError("The chat template rewrote history while calculating a non-user token delta") + + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + return new_tokens, [0] * len(new_tokens) diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py index 6e2043f9b..2fe5118e8 100644 --- a/examples/tau-bench/trainable_agents.py +++ b/examples/tau-bench/trainable_agents.py @@ -15,6 +15,7 @@ from tau_bench.agents.tool_calling_agent import RESPOND_ACTION_NAME from tau_bench.envs import get_env from tau_bench.types import Action, RunConfig +from token_delta import get_token_delta from vime.rollout.vllm_rollout import ( GenerateState, @@ -406,7 +407,6 @@ def sampling_params_for_turn() -> dict | None: return params try: - pending_obs_offset: int | None = None rendered_body = await safe_render() if rendered_body is None: return _mark_truncated() @@ -426,18 +426,6 @@ def sampling_params_for_turn() -> dict | None: return _mark_truncated() for turn_idx in range(args.max_turns): - input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) - - if pending_obs_offset is not None: - obs_tokens = input_ids[pending_obs_offset:] - remaining = remaining_budget() - if remaining is not None and len(obs_tokens) > remaining: - append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0)) - sample.status = Sample.Status.TRUNCATED - break - append_response_window(obs_tokens, [0] * len(obs_tokens)) - pending_obs_offset = None - current_sampling_params = sampling_params_for_turn() if current_sampling_params is None: sample.status = Sample.Status.TRUNCATED @@ -506,12 +494,25 @@ def sampling_params_for_turn() -> dict | None: train_logprobs.append(0.0) train_loss_mask.append(0) + has_previous_response = bool(response_tokens) response_tokens.extend(new_tokens) + messages.append({"role": "assistant", "content": response_text}) + assistant_delta_ids, assistant_delta_mask = get_token_delta( + state.tokenizer, + messages, + include_generation_prompt=has_previous_response, + ) + generation_prompt_length = next( + (index for index, mask in enumerate(assistant_delta_mask) if mask), + len(assistant_delta_mask), + ) + append_response_window( + assistant_delta_ids[:generation_prompt_length], + assistant_delta_mask[:generation_prompt_length], + ) append_response_window(train_tokens, train_loss_mask, train_logprobs) _maybe_apply_routed_experts(args, sample, choice) - messages.append({"role": "assistant", "content": response_text}) - if finish_reason == "length": sample.status = Sample.Status.TRUNCATED break @@ -527,15 +528,25 @@ def sampling_params_for_turn() -> dict | None: sample.status = Sample.Status.COMPLETED break + render_prefix_len = len(sample.tokens) next_user_message = env.format_observation(observation) messages.append(next_user_message) + obs_tokens, obs_loss_mask = get_token_delta(state.tokenizer, messages) + remaining = remaining_budget() + if remaining is not None and len(obs_tokens) > remaining: + append_response_window( + obs_tokens[: max(remaining, 0)], + obs_loss_mask[: max(remaining, 0)], + ) + sample.status = Sample.Status.TRUNCATED + break + append_response_window(obs_tokens, obs_loss_mask) if turn_idx + 1 >= args.max_turns: sample.reward = compute_process_reward(env, 0.0) sample.status = Sample.Status.TRUNCATED break - pending_obs_offset = len(input_ids) + len(train_tokens) max_ctx = args.rollout_max_context_len or 8192 if len(sample.tokens) >= max_ctx - 64: logger.info( @@ -548,10 +559,10 @@ def sampling_params_for_turn() -> dict | None: if rendered_body is None: return _mark_truncated() rendered_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) - is_prefix_stable = rendered_ids[:pending_obs_offset] == sample.tokens[:pending_obs_offset] + is_prefix_stable = rendered_ids[:render_prefix_len] == sample.tokens[:render_prefix_len] sample.metadata["multiturn_render"] = { "prefix_stable": is_prefix_stable, - "prefix_len": pending_obs_offset, + "prefix_len": render_prefix_len, "sample_len": len(sample.tokens), "rendered_len": len(rendered_ids), "turn": turn_idx + 1, diff --git a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh index 198f6f6ae..d23565eed 100644 --- a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh +++ b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh @@ -1,7 +1,7 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray @@ -28,6 +28,7 @@ source "/root/vime/scripts/models/qwen3-4B.sh" CKPT_ARGS=( --hf-checkpoint /root/Qwen3-4B + #--hf-checkpoint /root/Qwen3-4B-FP8 --ref-load /root/Qwen3-4B_torch_dist # --load /root/Qwen3-4B_vime/ --save /root/Qwen3-4B_vime/ @@ -153,4 +154,4 @@ ray job submit --address="http://127.0.0.1:8265" \ ${EVAL_ARGS[@]} \ ${VLLM_ARGS[@]} \ ${MISC_ARGS[@]} \ - ${CUSTOM_ARGS[@]} + ${CUSTOM_ARGS[@]} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 53e5766f9..52c20a24c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ build-backend = "setuptools.build_meta" profile = "black" # black-compatible line_length = 119 # should match black parameters ignore_whitespace = true # ignore whitespace for compatibility with the initial style -py_version = 310 # python 3.10 as a target version +py_version = 311 # python 3.11 as a target version sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"] default_section = "THIRDPARTY" extend_skip = ["setup.py", "docs/source/conf.py"] diff --git a/requirements.txt b/requirements.txt index b44dc96f1..0339d6495 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ accelerate anthropic +blake3 blobfile cloudpickle datasets @@ -16,9 +17,10 @@ pylatexenc pyyaml qwen_vl_utils # for VLM ray[default] -ring_flash_attn safetensors tensorboard transformers -vllm-router>=0.1.14 +vllm-router>=0.1.15 wandb +xxhash # disk delta weight sync checksum +zstandard diff --git a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh new file mode 100755 index 000000000..9e950d368 --- /dev/null +++ b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh @@ -0,0 +1,183 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/kimi-k2-thinking.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Kimi-K2-Thinking/ + --ref-load /root/Kimi-K2_thinking_torch_dist/ + --load /root/Kimi-K2-thinking_vime/ + --save /root/Kimi-K2-thinking_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type math + + --num-rollout 100 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 16384 + --rollout-temperature 0.8 + + # --global-batch-size 256 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 10 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 8 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 5 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + # --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group kimi-k2-thinking-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + + # dp attention + # --vllm-data-parallel-size 8 + + --vllm-enable-expert-parallel + + # enable deepep for vllm + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 +) + + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + + # use deepep for megatron + # --moe-enable-deepep + # --moe-token-dispatcher-type flex + --no-check-for-nan-in-loss-and-grad +) + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"NCCL_TIMEOUT_MS\":\"360000000\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"32\" + } +}" + + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 /personal/vime/vime/train.py \ + --actor-num-nodes 32 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/low_precision/run-moonlight-16B-A3B-int4.sh b/scripts/low_precision/run-moonlight-16B-A3B-int4.sh new file mode 100755 index 000000000..95365c0ff --- /dev/null +++ b/scripts/low_precision/run-moonlight-16B-A3B-int4.sh @@ -0,0 +1,166 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python +pkill -9 redis + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/moonlight.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Moonlight-16B-A3B-Instruct-INT4 + --ref-load /root/Moonlight-16B-A3B-Instruct-INT4_torch_dist + --load /root/Moonlight-16B-A3B_vime/ + --save /root/Moonlight-16B-A3B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 3000 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-temperature 0.8 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + # --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 4096 + --eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 4 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group moomlight-16B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 4 + --vllm-gpu-memory-utilization 0.7 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + # --attention-backend flash + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/low_precision/run-qwen3-235B-A22B-int4.sh b/scripts/low_precision/run-qwen3-235B-A22B-int4.sh new file mode 100755 index 000000000..df7e37ea0 --- /dev/null +++ b/scripts/low_precision/run-qwen3-235B-A22B-int4.sh @@ -0,0 +1,170 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/qwen3-235B-A22B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-235B-A22B-INT4/ + --ref-load /root/Qwen3-235B-A22B_torch_dist/ + --load /root/Qwen3-235B-A22B_vime/ + --save /root/Qwen3-235B-A22B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 300 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 0.8 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 10 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 2 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 22 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + # --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-235B-A22B-test + # --wandb-key ${WANDB_KEY} +) + + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + # --vllm-data-parallel-size 4 + --vllm-enable-expert-parallel + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + --no-check-for-nan-in-loss-and-grad + + # use deepep for megatron + # --moe-enable-deepep + # --moe-token-dispatcher-type flex +) + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"NCCL_TIMEOUT_MS\":\"360000000\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 8 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/low_precision/run-qwen3-30B-A3B-int4.sh b/scripts/low_precision/run-qwen3-30B-A3B-int4.sh new file mode 100755 index 000000000..03f98875d --- /dev/null +++ b/scripts/low_precision/run-qwen3-30B-A3B-int4.sh @@ -0,0 +1,165 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/qwen3-30B-A3B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-30B-A3B-INT4/ + --ref-load /root/Qwen3-30B-A3B_torch_dist/ + --load /root/Qwen3-30B-A3B_vime/ + --save /root/Qwen3-30B-A3B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 100 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 0.8 + + --global-batch-size 256 + --balance-data + # --debug-rollout-only +) + +EVAL_ARGS=( + --eval-interval 10 + --eval-prompt-data /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 16384 + --eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-30B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.7 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + # use deepep for megatron + # --moe-enable-deepep + # --moe-token-dispatcher-type flex +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} + diff --git a/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh b/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh new file mode 100755 index 000000000..1977fdd45 --- /dev/null +++ b/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh @@ -0,0 +1,180 @@ +#!/bin/bash + +# for rerun the task +# pkill -9 -f '[v]llm serve|VLL[M]::' +# sleep 3 +# ray stop --force +# pkill -9 ray +# pkill -9 python +# sleep 3 +# pkill -9 ray +# pkill -9 python +# pkill -9 redis + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/qwen3-30B-A3B.sh" + +# Base directory for checkpoints and related files (adjust if necessary) +BASE_DIR="/root" + +CKPT_ARGS=( + --hf-checkpoint "${BASE_DIR}/Qwen3-30B-A3B-FP8/" + --ref-load "${BASE_DIR}/Qwen3-30B-A3B_torch_dist/" + --load "${BASE_DIR}/Qwen3-30B-A3B_vime/" + --save "${BASE_DIR}/Qwen3-30B-A3B_vime/" + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data "${BASE_DIR}/dapo-math-17k.jsonl" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 200 + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 128 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime "${BASE_DIR}/aime-2024.jsonl" + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 1 + --expert-model-parallel-size 4 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 20480 + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex + + # fp8 + --transformer-impl transformer_engine + --bf16 + --fp8-format e4m3 + --fp8-recipe blockwise + # --fp8-param-gather +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-30B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.6 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + --vllm-enable-expert-parallel + # --use-rollout-routing-replay +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# Get Ray Head node info automatically +ip=$(ps aux | grep dashboard | grep -oP '(?<=--node-ip-address=)[0-9\.]+' | head -1) +port=$(ps aux | grep dashboard | grep -oP '(?<=dashboard-port=)\d+' | head -1) +export HEAD_NODE_ADDRESS="$ip" +export DASHBOARD_PORT="$port" +echo "Detected Ray Head IP: $HEAD_NODE_ADDRESS, Port: $DASHBOARD_PORT" + +export RAY_ADDRESS="http://${HEAD_NODE_ADDRESS}:${DASHBOARD_PORT}" + +# You should enable NVTE_FP8_BLOCK_SCALING_FP32_SCALES to use fp32 scales in fp8 training +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\": \"1\", + \"NCCL_TIMEOUT_MS\":\"36000000\" + } +}" + +ray job submit --address="${RAY_ADDRESS}" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 2 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/low_precision/run-qwen3-4b-fp8.sh b/scripts/low_precision/run-qwen3-4b-fp8.sh new file mode 100755 index 000000000..7b5caa3fc --- /dev/null +++ b/scripts/low_precision/run-qwen3-4b-fp8.sh @@ -0,0 +1,154 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../models/qwen3-4B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B-FP8 + #--hf-checkpoint /root/Qwen3-4B-FP8 + --ref-load /root/Qwen3-4B_torch_dist + --load /root/qwen3-4b_cp8_fp8 + --save /root/rl-model/qwen3-4b_cp8_fp8 + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/data/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-4B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.7 +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +# you should enable NVTE_FP8_BLOCK_SCALING_FP32_SCALES to use fp32 scales in fp8 training +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\": \"1\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ No newline at end of file diff --git a/scripts/models/glm4.7-30B-A3B-npu.sh b/scripts/models/glm4.7-30B-A3B-npu.sh deleted file mode 100644 index 9cc66c850..000000000 --- a/scripts/models/glm4.7-30B-A3B-npu.sh +++ /dev/null @@ -1,51 +0,0 @@ -MOE_ROUTED_EXPERTS=64 -MOE_ACTIVE_ROUTED_EXPERTS=4 -MOE_SHARED_EXPERTS=1 - -NHIDDEN=2048 -MOE_FFN_HIDDEN=1536 -MOE_SHARED_EXPERT_INTERMEDIATE_SIZE=$((MOE_FFN_HIDDEN * MOE_SHARED_EXPERTS)) -FFN_HIDDEN=10240 -N_DENSE_LAYERS=1 -N_MOE_LAYERS=46 -NHEADS=20 - -MODEL_ARGS=( - --moe-layer-freq "[0]*$N_DENSE_LAYERS+[1]*$N_MOE_LAYERS" - --num-experts $MOE_ROUTED_EXPERTS - --moe-shared-expert-intermediate-size $MOE_SHARED_EXPERT_INTERMEDIATE_SIZE - --moe-router-topk $MOE_ACTIVE_ROUTED_EXPERTS - --moe-grouped-gemm - --moe-ffn-hidden-size $MOE_FFN_HIDDEN - --moe-router-score-function sigmoid - --moe-router-pre-softmax - --moe-router-enable-expert-bias - --moe-router-bias-update-rate 0 - --moe-router-load-balancing-type seq_aux_loss - --moe-router-topk-scaling-factor 1.8 - --moe-aux-loss-coeff 0 - --moe-router-dtype fp32 - --make-vocab-size-divisible-by 64 - --num-layers $((N_DENSE_LAYERS + N_MOE_LAYERS)) - --hidden-size $NHIDDEN - --ffn-hidden-size $FFN_HIDDEN - --num-attention-heads $NHEADS - --disable-bias-linear - --add-qkv-bias - --swiglu - --untie-embeddings-and-output-weights - --position-embedding-type rope - --no-position-embedding - --normalization RMSNorm - --qk-layernorm - --multi-latent-attention - --q-lora-rank 768 - --kv-lora-rank 512 - --qk-head-dim 192 - --v-head-dim 256 - --kv-channels 192 - --qk-pos-emb-head-dim 64 - --vocab-size 154880 - --rotary-base 1000000 - --no-rope-fusion -) diff --git a/scripts/models/glm5.2-744B-A40B.sh b/scripts/models/glm5.2-744B-A40B.sh new file mode 100644 index 000000000..b86c81d9e --- /dev/null +++ b/scripts/models/glm5.2-744B-A40B.sh @@ -0,0 +1,62 @@ +MOE_ROUTED_EXPERTS=256 +MOE_ACTIVE_ROUTED_EXPERTS=8 +MOE_SHARED_EXPERTS=1 + +NHIDDEN=6144 +MOE_FFN_HIDDEN=2048 +MOE_SHARED_EXPERT_INTERMEDIATE_SIZE=$(($MOE_FFN_HIDDEN * $MOE_SHARED_EXPERTS)) +FFN_HIDDEN=12288 +N_DENSE_LAYERS=3 +N_MOE_LAYERS=75 +NHEADS=64 + +# GLM-5.2 744B-A40B with DSA *cross-layer index sharing*. +# Only the computing layers (layer 1,2,3,7,11,...,75 in Megatron 1-indexing) +# carry indexer weights and compute the sparse top-k; the remaining layers reuse +# the most recent computing layer's indices. The schedule (index_topk_freq=4, +# index_skip_topk_offset=3) is read from the HF config by the shared glm5 provider +# (cross-layer sharing activates automatically when index_topk_freq > 1). +MODEL_ARGS=( + --spec "vime_plugins.models.glm5.glm5" "get_glm5_spec" + --moe-layer-freq [0]*$N_DENSE_LAYERS+[1]*$N_MOE_LAYERS + --num-experts $MOE_ROUTED_EXPERTS + --moe-shared-expert-intermediate-size $MOE_SHARED_EXPERT_INTERMEDIATE_SIZE + --moe-router-topk $MOE_ACTIVE_ROUTED_EXPERTS + --moe-grouped-gemm + --moe-permute-fusion + --moe-ffn-hidden-size $MOE_FFN_HIDDEN + --moe-router-score-function sigmoid + --moe-router-pre-softmax + --moe-router-enable-expert-bias + --moe-router-bias-update-rate 0 + --moe-router-load-balancing-type seq_aux_loss + --moe-router-topk-scaling-factor 2.5 + --moe-aux-loss-coeff 0 + --moe-router-dtype fp32 + --make-vocab-size-divisible-by 16 + --num-layers $((N_DENSE_LAYERS + N_MOE_LAYERS)) + --hidden-size $NHIDDEN + --ffn-hidden-size $FFN_HIDDEN + --num-attention-heads $NHEADS + --disable-bias-linear + --swiglu + --untie-embeddings-and-output-weights + --position-embedding-type rope + --no-position-embedding + --normalization RMSNorm + --qk-layernorm + --multi-latent-attention + --q-lora-rank 2048 + --kv-lora-rank 512 + --qk-head-dim 192 + --v-head-dim 256 + --kv-channels 192 + --qk-pos-emb-head-dim 64 + --vocab-size 154880 + --rotary-base 8000000 + --enable-experimental + + # DSA + context parallel uses the sequential allgather-CP layout (not zigzag); + # the index-share provider gathers index_k/kv across the CP group to match. + --allgather-cp +) diff --git a/scripts/models/gpt-oss-20B.sh b/scripts/models/gpt-oss-20B.sh deleted file mode 100755 index 73bc6d57c..000000000 --- a/scripts/models/gpt-oss-20B.sh +++ /dev/null @@ -1,55 +0,0 @@ -#!/bin/bash - -# GPT-OSS 20B model configuration -# Based on openai/gpt-oss-20b -# 24 layers, 2880 hidden, 64 heads (8 kv), 32 experts top-4, all MoE -# Features: learnable softmax, SWA (window=128, skip_freq=2), quick GeGLU - -NLAYERS=24 - -MODEL_ARGS=( - --spec "vime_plugins.models.gpt_oss" "get_gpt_oss_spec" - --num-layers ${NLAYERS} - --hidden-size 2880 - --ffn-hidden-size 2880 - --num-attention-heads 64 - --group-query-attention - --num-query-groups 8 - --kv-channels 64 - --seq-length 4096 - --max-position-embeddings 131072 - --padded-vocab-size 201088 - --make-vocab-size-divisible-by 128 - --tokenizer-type HuggingFaceTokenizer - --bf16 - --normalization RMSNorm - --untie-embeddings-and-output-weights - --no-masked-softmax-fusion - --no-rope-fusion - --no-bias-gelu-fusion - --no-bias-dropout-fusion - --use-mcore-models - --rotary-percent 1.0 - --rotary-base 150000 - --position-embedding-type rope - --use-rope-scaling - --rope-scaling-factor 32 - --sequence-parallel - # MoE - --num-experts 32 - --moe-ffn-hidden-size 2880 - --moe-router-topk 4 - --moe-router-dtype fp32 - --moe-router-score-function softmax - --moe-router-load-balancing-type none - --moe-aux-loss-coeff 0.0 - --moe-token-dispatcher-type alltoall - --moe-grouped-gemm - # GPT-OSS specific - --quick-geglu - --glu-linear-offset 1.0 - --softmax-type learnable - --window-attn-skip-freq 2 - --window-size 128,0 - --activation-func-clamp-value 7.0 -) diff --git a/scripts/models/qwen3.5-35B-A3B-vl.sh b/scripts/models/qwen3.5-35B-A3B-vl.sh new file mode 100644 index 000000000..bcd716aa3 --- /dev/null +++ b/scripts/models/qwen3.5-35B-A3B-vl.sh @@ -0,0 +1,4 @@ +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/qwen3.5-35B-A3B.sh" + +MODEL_ARGS[1]="vime_plugins.models.qwen3_5_vl" +MODEL_ARGS[2]="get_qwen3_5_vl_model_provider" diff --git a/scripts/models/qwen3.5-9B.sh b/scripts/models/qwen3.5-9B.sh new file mode 100644 index 000000000..f34749860 --- /dev/null +++ b/scripts/models/qwen3.5-9B.sh @@ -0,0 +1,28 @@ +MODEL_ARGS=( + --spec "vime_plugins.models.qwen3_5" "get_qwen3_5_spec" + + --disable-bias-linear + --qk-layernorm + --group-query-attention + --num-attention-heads 16 + --num-query-groups 4 + --kv-channels 256 + --num-layers 32 + --hidden-size 4096 + --ffn-hidden-size 12288 + --use-gated-attention + + --normalization RMSNorm + --apply-layernorm-1p + --position-embedding-type rope + --norm-epsilon 1e-6 + --rotary-percent 0.25 + --swiglu + --untie-embeddings-and-output-weights + --vocab-size 248320 + + --rotary-base 10000000 + + # Qwen3.5 specific + --attention-output-gate +) diff --git a/scripts/models/qwen3_omni_moe.sh b/scripts/models/qwen3_omni_moe.sh new file mode 100644 index 000000000..31da1406e --- /dev/null +++ b/scripts/models/qwen3_omni_moe.sh @@ -0,0 +1,6 @@ +source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/qwen3-30B-A3B.sh" + +MODEL_ARGS+=( + --spec vime_plugins.models.qwen3_omni_moe get_qwen3_omni_model_provider + --use-qwen-vl +) diff --git a/scripts/run-deepseek-r1.sh b/scripts/run-deepseek-r1.sh new file mode 100755 index 000000000..76cb982c2 --- /dev/null +++ b/scripts/run-deepseek-r1.sh @@ -0,0 +1,171 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/deepseek-v3.sh" + +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/DeepSeek-R1/ + #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ + --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/ + --load $BASE_DIR/DeepSeek-R1_vime/ + --save $BASE_DIR/DeepSeek-R1_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 32768 + --rollout-temperature 1 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime $BASE_DIR/rl_data/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 32768 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 13 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group deepseek-r1-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.7 + --vllm-enable-expert-parallel + + # dp attention + --vllm-data-parallel-size 8 + + # enable deepep for vllm + + # mtp + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) + +# launch the master node of ray in container +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}", + "MASTER_ADDR": "${MASTER_ADDR}", + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + } + }' \ + -- python3 train.py \ + --actor-num-nodes 16 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-glm4-9B.sh b/scripts/run-glm4-9B.sh new file mode 100755 index 000000000..99c2bdc8e --- /dev/null +++ b/scripts/run-glm4-9B.sh @@ -0,0 +1,150 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/glm4-9B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/GLM-Z1-9B-0414/ + --ref-load /root/GLM-Z1-9B-0414_torch_dist + --load /root/GLM-Z1-9B-0414_vime/ + --save /root/GLM-Z1-9B-0414_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 2 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 4608 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-4B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-glm4.7-30B-A3B-npu.sh b/scripts/run-glm4.7-30B-A3B-npu.sh index 1b94be2bc..0ac6d7125 100644 --- a/scripts/run-glm4.7-30B-A3B-npu.sh +++ b/scripts/run-glm4.7-30B-A3B-npu.sh @@ -24,24 +24,25 @@ export HYDRA_FULL_ERROR=1 export DISABLE_L2_CACHE=1 export VLLM_ASCEND_ENABLE_NZ=0 export VLLM_USE_AOT_COMPILE=0 -export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}" +export PYTHONPATH="/root/Megatron-LM:/root/MegatronAdaptor:/root/TransformerEngineNPU:${PYTHONPATH:-}" unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4.7-30B-A3B-npu.sh" +source "${SCRIPT_DIR}/models/glm4.7-30B-A3B.sh" DATA_ROOT="${DATA_ROOT:-/root}" +MODEL_DIR="${DATA_ROOT}/weights/GLM-4.7-Flash" +DATASET_DIR="${DATA_ROOT}/datasets/dapo-math-17k" CKPT_ARGS=( - --hf-checkpoint ${DATA_ROOT}/weights/GLM-4.7-Flash/ - --load ${DATA_ROOT}/weights/GLM-4.7-Flash/ - --ref-load ${DATA_ROOT}/weights/GLM-4.7-Flash/ - --megatron-to-hf-mode bridge + --hf-checkpoint "${MODEL_DIR}" + --load "${MODEL_DIR}" + --ref-load "${MODEL_DIR}" ) ROLLOUT_ARGS=( - --prompt-data ${DATA_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data "${DATASET_DIR}/dapo-math-17k.jsonl" --input-key prompt --label-key label --apply-chat-template @@ -81,12 +82,6 @@ PERF_ARGS=( --seq-length 24576 ) -MTP_ARGS=( - --mtp-num-layers 1 - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) - GRPO_ARGS=( --advantage-estimator grpo --use-kl-loss @@ -111,13 +106,16 @@ OPTIMIZER_ARGS=( VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 --vllm-gpu-memory-utilization 0.7 + --vllm-enable-expert-parallel --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":1}' ) MISC_ARGS=( + # Match GLM's unscaled RoPE without changing main's shared model script. + --rope-type rope --attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 @@ -130,7 +128,7 @@ MISC_ARGS=( # launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --num-gpus 0 --resources '{"NPU": 16}' --node-ip-address ${MASTER_ADDR} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 ray job submit --address="http://127.0.0.1:8265" \ -- python3 train.py \ @@ -145,5 +143,4 @@ ray job submit --address="http://127.0.0.1:8265" \ "${PERF_ARGS[@]}" \ "${EVAL_ARGS[@]}" \ "${VLLM_ARGS[@]}" \ - "${MISC_ARGS[@]}" \ - "${MTP_ARGS[@]}" + "${MISC_ARGS[@]}" diff --git a/scripts/run-glm4.7-30B-A3B.sh b/scripts/run-glm4.7-30B-A3B.sh index b1cbc0fc1..75c2a3d03 100644 --- a/scripts/run-glm4.7-30B-A3B.sh +++ b/scripts/run-glm4.7-30B-A3B.sh @@ -1,7 +1,7 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray @@ -116,10 +116,11 @@ VLLM_ARGS=( --rollout-num-gpus-per-engine 8 --vllm-gpu-memory-utilization 0.8 --vllm-data-parallel-size 8 + --vllm-enable-expert-parallel # mtp - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - --vllm-max-cudagraph-capture-size 64 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' + --vllm-max-cudagraph-capture-size 320 --vllm-max-num-seqs 512 ) @@ -146,6 +147,7 @@ RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" } }" diff --git a/scripts/run-glm4.7-355B-A32B.sh b/scripts/run-glm4.7-355B-A32B.sh index 8dfe2ebf0..02326caba 100644 --- a/scripts/run-glm4.7-355B-A32B.sh +++ b/scripts/run-glm4.7-355B-A32B.sh @@ -1,7 +1,7 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray @@ -117,7 +117,7 @@ VLLM_ARGS=( # mtp --vllm-enable-expert-parallel - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' ) MISC_ARGS=( @@ -156,7 +156,7 @@ if [ -n "${HOSTFILE}" ]; then fi echo "Starting Ray worker on ${WORKER_IP}" ssh root@"${WORKER_IP}" \ - "pkill -9 -f 'vllm serve' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & done wait fi @@ -170,6 +170,7 @@ RUNTIME_ENV_JSON=$(cat </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/glm5-744B-A40B.sh" + +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/GLM-5 + --ref-load $BASE_DIR/GLM-5_torch_dist/ + --load $BASE_DIR/GLM-5_vime/ + --save $BASE_DIR/GLM-5_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3000 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 32768 + --rollout-temperature 1 + + --global-batch-size 64 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 4 + --decoder-last-pipeline-num-layers 18 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --context-parallel-size 2 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 + --data-pad-size-multiplier 4096 + --log-probs-chunk-size 1024 +) + +GRPO_ARGS=( + --advantage-estimator grpo + #--use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group glm5-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 64 + --vllm-gpu-memory-utilization 0.70 + --vllm-enable-expert-parallel + --vllm-data-parallel-size 64 + + + --prefill-num-servers 1 + + # mtp + + # dsa + --vllm-attention-backend nsa + --vllm-max-cudagraph-capture-size 40 + + --vllm-max-num-seqs 512 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' + +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"INDEXER_ROPE_NEOX_STYLE\": \"0\", + \"NVSHMEM_DISABLE_NCCL\": \"1\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ +-- python3 train.py \ + --actor-num-nodes 32 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + --update-weight-buffer-size $(( 1024 * 1024 * 1024 * 2 )) \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-glm5.2-744B-A40B.sh b/scripts/run-glm5.2-744B-A40B.sh new file mode 100644 index 000000000..28bffba2f --- /dev/null +++ b/scripts/run-glm5.2-744B-A40B.sh @@ -0,0 +1,259 @@ +#!/bin/bash + +# GLM-5.2 744B-A40B RL training on 32 nodes / 256 H100 GPUs with PD disaggregation. + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +export PYTHONUNBUFFERED=1 +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/glm5.2-744B-A40B.sh" + +if [ -z "${BASE_DIR:-}" ]; then + echo "BASE_DIR is not set. Please set it to a shared path visible from every node." + exit 1 +fi + +SOCKET_IFNAME=${SOCKET_IFNAME:-eth0} + +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/GLM-5.2-FP8 + --ref-load $BASE_DIR/GLM-5.2_torch_dist + --load $BASE_DIR/GLM-5.2_vime + --save $BASE_DIR/GLM-5.2_vime + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3000 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 65536 + --rollout-temperature 1.0 + + --global-batch-size 64 +) + +# TP=4, PP=8, CP=8 consumes all 256 GPUs (32 nodes) for one training group; DP=1. +# Experts use EP=32: expert_tp(1) * ep(32) * pp(8) = 256 = world_size (expert_dp=1). +# +# DSA cross-layer index sharing requires every pipeline stage to START on a +# "computing" layer (index_topk_freq=4, index_skip_topk_offset=3 -> computing +# layers are 1,2,3,7,11,...,75). A uniform 78/8 split would start stages on skip +# layers and fail. We instead use first=14, last=16, leaving 6 middle stages of +# (78-14-16)/6 = 8 layers each. Stage starts land on global layers +# 1,15,23,31,39,47,55,63 -- all computing layers. +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 8 + --decoder-first-pipeline-num-layers 14 + --decoder-last-pipeline-num-layers 16 + --context-parallel-size 8 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 + --data-pad-size-multiplier 1024 + --log-probs-chunk-size 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + + --use-tis + --tis-clip-low 0.5 + --tis-clip 2.0 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group glm5.2-744B-A40B +) + +VLLM_CONFIG_FILE=$(mktemp /tmp/vllm_glm52_744B_A40B_XXXXXX.yaml) +# PD disaggregation: 1 prefill engine (64 GPU) + 3 decode engines (192 GPU) = 256. +# Each engine spans 64 GPUs (EP=64, within DeepEP's supported rank set). Prefill +# uses the high-throughput DeepEP backend; decode uses the low-latency backend. +cat > "${VLLM_CONFIG_FILE}" </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/kimi-k2.sh" + +CKPT_ARGS=( + --hf-checkpoint $BASE_DIR/Kimi-K2-Instruct/ + # --hf-checkpoint $BASE_DIR/Kimi-K2-bf16/ + --ref-load $BASE_DIR/Kimi-K2_torch_dist/ + --load $BASE_DIR/Kimi-K2_vime/ + --save $BASE_DIR/Kimi-K2_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type math + + --num-rollout 100 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 32768 + --rollout-temperature 1 + + # --global-batch-size 1024 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime $BASE_DIR/rl_data/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 32768 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 8 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 5 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group kimi-k2-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 16 + --vllm-gpu-memory-utilization 0.7 + + # dp attention + --vllm-data-parallel-size 8 + + --vllm-enable-expert-parallel + + # enable deepep for vllm + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 +) + + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 32 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-kimi-k2-Thinking.sh b/scripts/run-kimi-k2-Thinking.sh new file mode 100755 index 000000000..915a83c46 --- /dev/null +++ b/scripts/run-kimi-k2-Thinking.sh @@ -0,0 +1,179 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/kimi-k2-thinking.sh" + +CKPT_ARGS=( + # --hf-checkpoint $BASE_DIR/Kimi-K2-Thinking-bf16/ + --hf-checkpoint $BASE_DIR/Kimi-K2-Thinking-fp8/ + --ref-load $BASE_DIR/Kimi-K2-Thinking_torch_dist/ + --load $BASE_DIR/Kimi-K2-Thinking_vime/ + --save $BASE_DIR/Kimi-K2-Thinking_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type math + + --num-rollout 100 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 16384 + --rollout-temperature 1 + + # --global-batch-size 1024 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime $BASE_DIR/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 8 + --context-parallel-size 4 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 5 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + # --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group kimi-k2-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 16 + --vllm-gpu-memory-utilization 0.7 + + # dp attention + --vllm-data-parallel-size 8 + + --vllm-enable-expert-parallel + + # enable deepep for vllm + + # make every dp rank has 128 concurrency + --vllm-server-concurrency 1024 +) + + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + + # use deepep for megatron + # --moe-enable-deepep + # --moe-token-dispatcher-type flex +) + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 32 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-mimo-7B-rl-eagle.sh b/scripts/run-mimo-7B-rl-eagle.sh new file mode 100755 index 000000000..0bcd3fca8 --- /dev/null +++ b/scripts/run-mimo-7B-rl-eagle.sh @@ -0,0 +1,163 @@ + +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/mimo-7B-rl.sh" + + CKPT_ARGS=( + --hf-checkpoint /root/MiMo-7B-RL + #--hf-checkpoint /root/Qwen3-4B-FP8 + --ref-load /root/MiMo-7B-RL_torch_dist + --load /root/MiMo-7B-RL-mtp_vime/ + --save /root/MiMo-7B-RL-mtp_vime/ + --save-interval 2000 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 1 + --eval-max-response-len 8192 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group mimo-7B-rl-test + # --wandb-key ${WANDB_API_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.7 + + # for speculative decoding + + # sometimes flashinfer has IMA bugs. Use fa3 as instead + --vllm-attention-backend fa3 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +SPEC_ARGS=( + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${SPEC_ARGS[@]} diff --git a/scripts/run-minimax-m2.sh b/scripts/run-minimax-m2.sh old mode 100644 new mode 100755 index 463d9dd8f..0eee7abcc --- a/scripts/run-minimax-m2.sh +++ b/scripts/run-minimax-m2.sh @@ -4,7 +4,7 @@ # ============================================================ # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray @@ -37,7 +37,6 @@ CKPT_ARGS=( --load ${BASE_DIR}/MiniMax-M2.5_vime/ --save ${BASE_DIR}/MiniMax-M2.5_vime/ --save-interval 20 - --megatron-to-hf-mode raw --model-name minimax_m2 ) diff --git a/scripts/run-moonlight-16B-A3B.sh b/scripts/run-moonlight-16B-A3B.sh new file mode 100755 index 000000000..2e62d004c --- /dev/null +++ b/scripts/run-moonlight-16B-A3B.sh @@ -0,0 +1,164 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python +pkill -9 redis + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/moonlight.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Moonlight-16B-A3B + --ref-load /root/Moonlight-16B-A3B_torch_dist + --load /root/Moonlight-16B-A3B_vime/ + --save /root/Moonlight-16B-A3B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 3000 + --rollout-batch-size 128 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-temperature 1 + + --over-sampling-batch-size 256 + --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std + + --num-steps-per-rollout 4 + # --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 4096 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group moomlight-16B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + # --attention-backend flash + + # use deepep for megatron + --moe-enable-deepep + --moe-token-dispatcher-type flex +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen2.5-0.5B-gb10-smoke.sh b/scripts/run-qwen2.5-0.5B-gb10-smoke.sh new file mode 100755 index 000000000..13b140fa3 --- /dev/null +++ b/scripts/run-qwen2.5-0.5B-gb10-smoke.sh @@ -0,0 +1,113 @@ +#!/bin/bash +# Minimal GRPO smoke test for vime on NVIDIA DGX Spark (GB10, single GPU). +# Goal: exercise the full rollout → reward → policy-update loop for one tiny +# step and exit cleanly. Used only to validate the GB10 port; not a training +# recipe. +# +# Prerequisites: +# - /root/Qwen2.5-0.5B-Instruct (HF checkpoint) +# - /root/Qwen2.5-0.5B-Instruct_torch_dist (from tools/convert_hf_to_torch_dist.py) +# - /root/dapo-math-17k/dapo-math-17k.jsonl + +set -ex + +# clean any leftover ray/vllm +pkill -9 -f '[v]llm serve|VLL[M]::' 2>/dev/null || true +ray stop --force 2>/dev/null || true +pkill -9 ray python 2>/dev/null || true +sleep 2 + +export PYTHONUNBUFFERED=1 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen2.5-0.5B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/ + --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ + --save /tmp/vime_smoke_save/ + --save-interval 9999 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + + --num-rollout 1 + --rollout-batch-size 2 + --n-samples-per-prompt 2 + --num-steps-per-rollout 1 + --global-batch-size 4 + + --rollout-max-response-len 256 + --rollout-temperature 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 2048 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization 0.4 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +ray start --head --node-ip-address 127.0.0.1 --num-gpus 1 --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/src/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1" + } + }' \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --colocate \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" diff --git a/scripts/run-qwen3-235B-A22B-sft.sh b/scripts/run-qwen3-235B-A22B-sft.sh new file mode 100755 index 000000000..62b7eae73 --- /dev/null +++ b/scripts/run-qwen3-235B-A22B-sft.sh @@ -0,0 +1,151 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# if base folder not set raise error +if [ -z "${BASE_FOLDER}" ]; then + echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." + exit 1 +fi + +if [ -z "${MASTER_ADDR}" ]; then + echo "MASTER_ADDR is not set. Please set it to the master node address." + exit 1 +fi + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-235B-A22B.sh" + +CKPT_ARGS=( + --hf-checkpoint ${BASE_FOLDER}/Qwen3-235B-A22B + --ref-load ${BASE_FOLDER}/Qwen3-235B-A22B_torch_dist + --load ${BASE_FOLDER}/Qwen3-235B-A22B_vime/ + --save ${BASE_FOLDER}/Qwen3-235B-A22B_vime/ + --save-interval 1000 +) + +SFT_ARGS=( + --rollout-function-path vime.rollout.sft_rollout.generate_rollout + --prompt-data ${BASE_FOLDER}/openhermes2_5.parquet + --input-key messages + # --apply-chat-template + --rollout-shuffle + --num-epoch 3 + --rollout-batch-size 128 + --global-batch-size 128 + + --loss-type sft_loss + --calculate-per-token-loss + --disable-compute-advantages-and-returns + --debug-train-only +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 32 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --lr-decay-style cosine + --min-lr 1e-6 + --lr-warmup-fraction 0.1 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-235B-sft +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do + if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & +done +wait + + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 4 \ + --actor-num-gpus-per-node 8 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${SFT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-235B-A22B.sh b/scripts/run-qwen3-235B-A22B.sh new file mode 100755 index 000000000..fedd91e3f --- /dev/null +++ b/scripts/run-qwen3-235B-A22B.sh @@ -0,0 +1,183 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# if base folder not set raise error +if [ -z "${BASE_FOLDER}" ]; then + echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." + exit 1 +fi + +if [ -z "${MASTER_ADDR}" ]; then + echo "MASTER_ADDR is not set. Please set it to the master node address." + exit 1 +fi + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-235B-A22B.sh" + +CKPT_ARGS=( + --hf-checkpoint ${BASE_FOLDER}/Qwen3-235B-A22B-FP8 + --ref-load ${BASE_FOLDER}/Qwen3-235B-A22B_torch_dist + --load ${BASE_FOLDER}/Qwen3-235B-A22B_vime/ + --save ${BASE_FOLDER}/Qwen3-235B-A22B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data ${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3000 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 64 + --balance-data +) + +EVAL_ARGS=( + #--eval-interval 20 + --eval-prompt-data aime ${BASE_FOLDER}/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 2 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + --decoder-last-pipeline-num-layers 22 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator gspo + #--use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 4e-4 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-235B-A22B +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 32 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 4 + --vllm-enable-expert-parallel + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do + if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & +done +wait + + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 8 \ + --actor-num-gpus-per-node 8 \ + --rollout-num-gpus 64 \ + --update-weight-buffer-size $(( 1024 * 1024 * 1024 * 4 )) \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-30B-A3B-npu.sh b/scripts/run-qwen3-30B-A3B-npu.sh index f7ae64543..87a717dff 100644 --- a/scripts/run-qwen3-30B-A3B-npu.sh +++ b/scripts/run-qwen3-30B-A3B-npu.sh @@ -2,7 +2,6 @@ # for rerun the task pkill -9 -f '[v]llm serve|VLL[M]::' -pkill -9 -f VLLM sleep 3 ray stop --force pkill -9 ray @@ -37,7 +36,6 @@ CKPT_ARGS=( --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-30B-A3B/ --load ${DATA_ROOT}/weights/Qwen3-30B-A3B/ --ref-load ${DATA_ROOT}/weights/Qwen3-30B-A3B/ - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -104,6 +102,7 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 --vllm-gpu-memory-utilization 0.7 --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) @@ -120,7 +119,7 @@ MISC_ARGS=( --no-gradient-accumulation-fusion ) -ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --num-gpus 0 --resources '{"NPU": 16}' --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 ray job submit --address="http://127.0.0.1:8265" \ -- python3 train.py \ diff --git a/scripts/run-qwen3-30B-A3B.sh b/scripts/run-qwen3-30B-A3B.sh new file mode 100644 index 000000000..745d52446 --- /dev/null +++ b/scripts/run-qwen3-30B-A3B.sh @@ -0,0 +1,156 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python +pkill -9 redis + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-30B-A3B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-30B-A3B + #--hf-checkpoint /root/Qwen3-30B-A3B-FP8 + --ref-load /root/Qwen3-30B-A3B_torch_dist + --load /root/Qwen3-30B-A3B_vime/ + --save /root/Qwen3-30B-A3B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 20480 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-30B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-32B.sh b/scripts/run-qwen3-32B.sh new file mode 100755 index 000000000..aa58ee204 --- /dev/null +++ b/scripts/run-qwen3-32B.sh @@ -0,0 +1,154 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-32B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-32B + --ref-load /root/Qwen3-32B_torch_dist/ + --load /root/Qwen3-32B_vime + --save /root/Qwen3-32B_vime + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 5 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 20480 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-30B-A3B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.7 + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-4B-base-sft.sh b/scripts/run-qwen3-4B-base-sft.sh new file mode 100755 index 000000000..849120090 --- /dev/null +++ b/scripts/run-qwen3-4B-base-sft.sh @@ -0,0 +1,127 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B-Base/ + --ref-load /root/Qwen3-4B-Base_torch_dist + --load /root/Qwen3-4B-Base_vime/ + --save /root/Qwen3-4B-Base_vime/ + --save-interval 1000 +) + +SFT_ARGS=( + --rollout-function-path vime.rollout.sft_rollout.generate_rollout + --prompt-data /root/openhermes2_5.parquet + --input-key messages + # --apply-chat-template + --rollout-shuffle + --num-epoch 3 + --rollout-batch-size 128 + --global-batch-size 128 + + --loss-type sft_loss + --calculate-per-token-loss + --disable-compute-advantages-and-returns + --debug-train-only +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --lr-decay-style cosine + --min-lr 1e-6 + --lr-warmup-fraction 0.1 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.95 +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-4B-base-sft + # --wandb-key ${WANDB_KEY} +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${SFT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-4B-npu.sh b/scripts/run-qwen3-4B-npu.sh index 73b672866..fe5fd4d3c 100644 --- a/scripts/run-qwen3-4B-npu.sh +++ b/scripts/run-qwen3-4B-npu.sh @@ -37,7 +37,6 @@ CKPT_ARGS=( --hf-checkpoint ${DATA_ROOT}/models/Qwen3-4B/ --load ${DATA_ROOT}/models/Qwen3-4B/ --ref-load ${DATA_ROOT}/models/Qwen3-4B/ - --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( @@ -67,7 +66,6 @@ PERF_ARGS=( --recompute-num-layers 1 --use-dynamic-batch-size --max-tokens-per-gpu 8192 - --megatron-to-hf-mode bridge ) GRPO_ARGS=( @@ -93,6 +91,7 @@ OPTIMIZER_ARGS=( ) VLLM_ARGS=( + --vllm-additional-config '{"weight_nz_mode":0}' --rollout-num-gpus-per-engine 4 --vllm-gpu-memory-utilization 0.6 ) @@ -107,7 +106,7 @@ MISC_ARGS=( --use-flash-attn ) -ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --num-gpus 0 --resources '{"NPU": 8}' --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 ray job submit --address="http://127.0.0.1:8265" \ -- python3 train.py \ diff --git a/scripts/run-qwen3-4B.sh b/scripts/run-qwen3-4B.sh new file mode 100644 index 000000000..200e3e687 --- /dev/null +++ b/scripts/run-qwen3-4B.sh @@ -0,0 +1,161 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +if command -v nvidia-smi >/dev/null 2>&1; then + DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ') +else + DETECTED_GPUS=0 +fi +NUM_GPUS=${NUM_GPUS:-${DETECTED_GPUS}} +if [ -z "$NUM_GPUS" ] || [ "$NUM_GPUS" -le 0 ]; then + NUM_GPUS=8 +fi +echo "NUM_GPUS: $NUM_GPUS" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-4B + #--hf-checkpoint /root/Qwen3-4B-FP8 + --ref-load /root/Qwen3-4B_torch_dist + --load /root/Qwen3-4B_vime/ + --save /root/Qwen3-4B_vime/ + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-4B-test + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.7 +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node ${NUM_GPUS} \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-8B-amd.sh b/scripts/run-qwen3-8B-amd.sh new file mode 100644 index 000000000..eb9b033a0 --- /dev/null +++ b/scripts/run-qwen3-8B-amd.sh @@ -0,0 +1,162 @@ +#!/bin/bash +# Colocate GRPO for Qwen3-8B on ROCm (gfx950 / MI350). Short validation run by +# default; override the knobs below to scale up. + +# Clean leftovers from a previous run (vLLM orphans procs named VLLM::*). +ray stop --force +pkill -9 -f "VLLM::" +pkill -9 -f "EngineCore" +pkill -9 -f "ray::" +pkill -9 -f "raylet|gcs_server|ray/dashboard|default_worker|log_monitor|runtime_env_agent|autoscaler" +pkill -9 -f "train.py" +sleep 3 +pkill -9 -f "VLLM::" + +set -ex + +export PYTHONUNBUFFERED=1 + +# Clear baked NVTE_* so Megatron sets the attn backend from --attention-backend flash. +unset NVTE_FUSED_ATTN NVTE_FLASH_ATTN NVTE_UNFUSED_ATTN + +# ----- ROCm GPU visibility: single TP=2 engine; keep HIP == CUDA visibility ----- +export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES:-1} +export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES:-1} +export HIP_VISIBLE_DEVICES=${VISIBLE_GPUS:-${HIP_VISIBLE_DEVICES:-6,7}} +export CUDA_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES}" + +IFS=',' read -r -a _visible_gpu_ids <<< "${CUDA_VISIBLE_DEVICES}" +NUM_GPUS=${NUM_GPUS:-${#_visible_gpu_ids[@]}} +HAS_NVLINK=0 # AMD: no NVLink; disable NCCL NVLS +echo "HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES NUM_GPUS=$NUM_GPUS HAS_NVLINK=$HAS_NVLINK" + +# ----- short-run knobs (override to scale up to a real job) ----- +NUM_ROLLOUT=${NUM_ROLLOUT:-3} +ROLLOUT_BATCH_SIZE=${ROLLOUT_BATCH_SIZE:-16} +N_SAMPLES_PER_PROMPT=${N_SAMPLES_PER_PROMPT:-8} +MAX_RESPONSE_LEN=${MAX_RESPONSE_LEN:-4096} +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-128} +EVAL_INTERVAL=${EVAL_INTERVAL:-100000} # effectively off for the short run + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_ROOT="$(cd -- "${SCRIPT_DIR}/.." &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-8B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-8B + --ref-load /root/Qwen3-8B_torch_dist + --load /root/Qwen3-8B_vime/ + --save /root/Qwen3-8B_vime/ + --save-interval 100000 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size ${ROLLOUT_BATCH_SIZE} + --n-samples-per-prompt ${N_SAMPLES_PER_PROMPT} + --rollout-max-response-len ${MAX_RESPONSE_LEN} + --rollout-temperature 1 + + --global-batch-size ${GLOBAL_BATCH_SIZE} + --balance-data +) + +# Eval disabled (no --eval-prompt-data). Add it back to measure accuracy. +EVAL_ARGS=() + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 4096 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + --use-wandb + --wandb-project "${WANDB_PROJECT:-vime-qwen3-8B-rocm}" + --wandb-group "${WANDB_GROUP:-qwen3-8B-grpo-rocm}" + --wandb-key "${WANDB_API_KEY:-${wandb_key}}" + --wandb-mode online +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + # Modest KV reservation so Megatron's post-rollout onload doesn't OOM at TP=2. + --vllm-gpu-memory-utilization 0.4 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + # APEX not installed on this ROCm image -> disable fused grad accumulation. + --no-gradient-accumulation-fusion + # Keep train state resident: the colocate offload path leaks VRAM on ROCm gfx950 + # (torch_memory_saver VMM) and Qwen3-8B fits alongside vLLM, so offload is unneeded. + --no-offload-train +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --train-backend megatron \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node ${NUM_GPUS} \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-8B-async-amd.sh b/scripts/run-qwen3-8B-async-amd.sh new file mode 100644 index 000000000..edd37ef3c --- /dev/null +++ b/scripts/run-qwen3-8B-async-amd.sh @@ -0,0 +1,146 @@ +#!/bin/bash +# Async (non-colocated) GRPO for Qwen3-8B on ROCm (gfx950 / MI350): train_async.py, +# actor and rollout on disjoint GPUs, RCCL-broadcast weight sync (no colocate IPC). + +# for rerun the task +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +# Clear baked NVTE_* so Megatron sets the attn backend from --attention-backend flash. +unset NVTE_FUSED_ATTN NVTE_FLASH_ATTN NVTE_UNFUSED_ATTN + +# ----- ROCm GPU visibility (same recipe as the colocate script) ----- +export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES:-1} +export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES:-1} +export HIP_VISIBLE_DEVICES=${VISIBLE_GPUS:-${HIP_VISIBLE_DEVICES:-0,1,2,3}} +export CUDA_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES}" +IFS=',' read -r -a _vis <<< "${CUDA_VISIBLE_DEVICES}" +NUM_GPUS=${NUM_GPUS:-${#_vis[@]}} +HAS_NVLINK=0 + +# ----- async GPU split: actor pool + rollout pool are disjoint ----- +ACTOR_GPUS=${ACTOR_GPUS:-2} # actor: TP=2 on 2 GPUs (DP=1) +ROLLOUT_GPUS=${ROLLOUT_GPUS:-$((NUM_GPUS - ACTOR_GPUS))} # rollout: remaining GPUs +echo "NUM_GPUS=$NUM_GPUS ACTOR_GPUS=$ACTOR_GPUS ROLLOUT_GPUS=$ROLLOUT_GPUS HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES" + +# ----- short validation run (override to scale) ----- +NUM_ROLLOUT=${NUM_ROLLOUT:-5} + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_ROOT="$(cd -- "${SCRIPT_DIR}/.." &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-8B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-8B + --ref-load /root/Qwen3-8B_torch_dist + --load /root/Qwen3-8B_vime_async/ + --save /root/Qwen3-8B_vime_async/ + --save-interval 100000 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-temperature 1 + --global-batch-size 128 + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.85 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --no-gradient-accumulation-fusion +) + +WANDB_ARGS=( + --use-wandb + --wandb-project "${WANDB_PROJECT:-vime-qwen3-8B-rocm}" + --wandb-group "${WANDB_GROUP:-qwen3-8B-grpo-async-rocm}" + --wandb-key "${WANDB_API_KEY:-${wandb_key}}" + --wandb-mode online +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node ${ACTOR_GPUS} \ + --rollout-num-gpus ${ROLLOUT_GPUS} \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${WANDB_ARGS[@]} diff --git a/scripts/run-qwen3-next-80B-A3B.sh b/scripts/run-qwen3-next-80B-A3B.sh new file mode 100755 index 000000000..f15444bb9 --- /dev/null +++ b/scripts/run-qwen3-next-80B-A3B.sh @@ -0,0 +1,195 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +if [ -z "${BASE_FOLDER:-}" ]; then + echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." + exit 1 +fi + +MASTER_ADDR=${MASTER_ADDR:-} +if [ -z "${MASTER_ADDR}" ]; then + echo "MASTER_ADDR is not set. Please set it to the master node address." + exit 1 +fi + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +# unset proxy to avoid distributed startup issues +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-4} +ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8} +CP_SIZE=${CP_SIZE:-4} +SOCKET_IFNAME=${SOCKET_IFNAME:-eth0} + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-next-80B-A3B.sh" + +CKPT_ARGS=( + --hf-checkpoint "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking" + --ref-load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_torch_dist" + --load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_vime/" + --save "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_vime/" + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data "${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 32768 + --rollout-temperature 1 + --global-batch-size 64 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl" + --n-samples-per-eval-prompt 8 + --eval-max-response-len 32768 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 4 + --decoder-last-pipeline-num-layers 9 + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator gspo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 4e-4 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3-next-80B-A3B-32k + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 8 + --vllm-gpu-memory-utilization 0.8 + --vllm-enable-expert-parallel + --vllm-cudagraph-capture-sizes 5 10 20 40 $(seq 80 40 640) + + # mtp + + --vllm-max-num-seqs 256 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep +) + +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +HOSTFILE=${HOSTFILE:-} +if [ -n "${HOSTFILE}" ]; then + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & + done + wait +fi + +RUNTIME_ENV_JSON=$(cat </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3.5-27B.sh" + +CKPT_ARGS=( + --hf-checkpoint "${BASE_FOLDER}/Qwen3.5-27B" + --ref-load "${BASE_FOLDER}/Qwen3.5-27B_torch_dist/" + --load "${BASE_FOLDER}/Qwen3.5-27B_vime/" + --save "${BASE_FOLDER}/Qwen3.5-27B_vime/" + --save-interval 20 +) + +ROLLOUT_ARGS=( + --prompt-data "${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 32768 + --rollout-temperature 1.0 + --global-batch-size 64 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl" + --n-samples-per-eval-prompt 8 + --eval-max-response-len 32768 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 2 + --decoder-last-pipeline-num-layers 30 + --context-parallel-size "${CP_SIZE}" + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --calculate-per-token-loss + --max-tokens-per-gpu 8192 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3.5-27B-32k + # --wandb-key ${WANDB_KEY} +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.75 + --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":4}' +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +HOSTFILE=${HOSTFILE:-} +if [ -n "${HOSTFILE}" ]; then + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & + done + wait +fi + +RUNTIME_ENV_JSON=$(cat </dev/null || true -sleep 2 -npu-smi info 2>/dev/null | grep rayWorker | awk '{print $4}' | xargs -r kill -9 2>/dev/null || true -sleep 3 - -# Ray isolation: independent temp-dir, ports, and cleanup -export RAY_TMPDIR=/tmp/ray_vime_npu_qwen35_35b_a3b -export RAY_PORT=6379 -export RAY_DASHBOARD_PORT=8265 -export RAY_AGENT_PORT=52378 -unset RAY_ADDRESS RAY_REDIS_ADDRESS - -ray stop --force 2>/dev/null || true -rm -rf "${RAY_TMPDIR}" -sleep 2 - -project_name="vime" -exp_name="qwen35-35b-a3b-rl" -RAY_DATA_HOME=${RAY_DATA_HOME:-"/root/logs"} -start_time=$(date +"%Y%m%d_%H%M%S") -LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/${project_name}/${exp_name}"} -mkdir -p "${LOG_DIR}" -LOG_FILE="${LOG_DIR}/${start_time}.log" - -echo "Experiment Log will be saved to: ${LOG_FILE}" -VIME_DIR="/root/vime" - -# NPU environment -source /usr/local/Ascend/driver/bin/setenv.bash -source /usr/local/Ascend/ascend-toolkit/set_env.sh -source /usr/local/Ascend/nnal/atb/set_env.sh -export PYTHONPATH="${VIME_DIR}:/root/Megatron-LM:/vllm-workspace/vllm:/vllm-workspace/vllm-ascend:/root/Megatron-Bridge/src:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:${PYTHONPATH}" -export PYTHONUNBUFFERED=1 -export PYTORCH_NPU_ALLOC_CONF=expandable_segments:False -export CUDA_DEVICE_MAX_CONNECTIONS=1 -export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 -export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 -export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 -export HCCL_CONNECT_TIMEOUT=7200 -export HCCL_DETERMINISTIC=true -export VLLM_ASCEND_ENABLE_NZ=0 -export ASCEND_COREDUMP_SIGNAL=None -export ATB_MATMUL_SHUFFLE_K_ENABLE=0 -export ATB_LLM_LCOC_ENABLE=0 -export TASK_QUEUE_ENABLE=0 -export RAY_DISABLE_SIGINT_OVERRIDE=1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export ASCEND_CUSTOM_OPP_PATH=/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer -export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:${LD_LIBRARY_PATH} -export VLLM_DISABLE_COMPILE_CACHE=1 -export TRANSFORMERS_VERBOSITY=error -export RUST_LOG=vllm_router_rs=warn - -NUM_NPUS=16 -source "${VIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint /path/to/Qwen3.5-35B-A3B - --ref-load /path/to/Qwen3.5-35B-A3B - --load /path/to/Qwen3.5-35B-A3B_vime_npu/ - --save /path/to/Qwen3.5-35B-A3B_vime_npu/ - --save-interval 20 - --no-load-optim - --megatron-to-hf-mode bridge -) - -ROLLOUT_ARGS=( - --prompt-data /path/to/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type deepscaler - --num-rollout 200 - --rollout-batch-size 8 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - --global-batch-size 64 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 50 - --eval-prompt-data aime /path/to/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --micro-batch-size 1 - --qkv-format bshd - --max-tokens-per-gpu 9216 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.85 - --vllm-enable-sleep-mode - --vllm-weight-sync-mode native -) - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash - --use-flash-attn -) - -# launch the master node of ray in container -unset https_proxy http_proxy proxy -ray start --head \ - --temp-dir="${RAY_TMPDIR}" \ - --port="${RAY_PORT}" \ - --dashboard-port="${RAY_DASHBOARD_PORT}" \ - --dashboard-agent-listen-port="${RAY_AGENT_PORT}" \ - --node-ip-address 127.0.0.1 \ - --num-gpus 0 \ - --resources "{\"NPU\": $NUM_NPUS}" \ - --disable-usage-stats \ - --dashboard-host=0.0.0.0 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON=$(cat << 'EOF' -{ - "env_vars": { - "PYTHONPATH": "${VIME_DIR}:/root/Megatron-LM:/vllm-workspace/vllm:/vllm-workspace/vllm-ascend:/root/Megatron-Bridge/src:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages", - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050", - "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050", - "HCCL_CONNECT_TIMEOUT": "7200", - "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False", - "VLLM_DISABLE_COMPILE_CACHE": "1", - "TRANSFORMERS_VERBOSITY": "error", - "RUST_LOG": "vllm_router_rs=warn", - "ASCEND_CUSTOM_OPP_PATH": "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer", - "LD_LIBRARY_PATH": "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib" - } -} -EOF -) - -ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - --working-dir="${VIME_DIR}" \ - -- python3 -u train.py \ - --train-backend megatron \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --rollout-num-gpus 8 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - 2>&1 | tee "${LOG_FILE}" \ No newline at end of file diff --git a/scripts/run-qwen3.5-35B-A3B-sft.sh b/scripts/run-qwen3.5-35B-A3B-sft.sh new file mode 100755 index 000000000..14f7a010a --- /dev/null +++ b/scripts/run-qwen3.5-35B-A3B-sft.sh @@ -0,0 +1,164 @@ +#!/bin/bash + +# for rerun the task +pkill -9 -f '[v]llm serve|VLL[M]::' +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# if base folder not set raise error +if [ -z "${BASE_FOLDER}" ]; then + echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." + exit 1 +fi + +if [ -z "${MASTER_ADDR}" ]; then + echo "MASTER_ADDR is not set. Please set it to the master node address." + exit 1 +fi +# export MASTER_ADDR="127.0.0.1" + +# will prevent ray from buffering stdout/stderr +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3.5-35B-A3B.sh" + +CKPT_ARGS=( + --hf-checkpoint ${BASE_FOLDER}/Qwen3.5-35B-A3B + --ref-load ${BASE_FOLDER}/Qwen3.5-35B-A3B_torch_dist + --load ${BASE_FOLDER}/Qwen3.5-35B-A3B_vime/ + --save ${BASE_FOLDER}/Qwen3.5-35B-A3B_vime/ + --save-interval 20 +) + +SFT_ARGS=( + --rollout-function-path vime.rollout.sft_rollout.generate_rollout + --prompt-data ${BASE_FOLDER}/openhermes2_5.parquet + --input-key messages + --rollout-shuffle + --num-epoch 3 + --rollout-batch-size 128 + --global-batch-size 128 + + --loss-type sft_loss + --loss-mask-type qwen3_5 + --calculate-per-token-loss + --disable-compute-advantages-and-returns + --debug-train-only +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --lr-decay-style cosine + --min-lr 1e-6 + --lr-warmup-fraction 0.1 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --use-distributed-optimizer + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project vime-dev + # --wandb-group qwen3.5-35B-sft +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + + --moe-token-dispatcher-type flex + --moe-enable-deepep +) + +SPEC_ARGS=( +# --mtp-num-layers 1 +# --enable-mtp-training +# --mtp-loss-scaling-factor 0.1 +) + +# launch the master node of ray in container +export no_proxy="127.0.0.1,${MASTER_ADDR}" +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do + if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then + continue + fi + echo "Starting Ray worker on ${WORKER_IP}" + ssh root@"${WORKER_IP}" \ + "pkill -9 -f '[v]llm serve|VLL[M]::' ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & +done +wait + + +# Build the runtime environment JSON with proper variable substitution +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NVSHMEM_DISABLE_NCCL\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"no_proxy\": \"${no_proxy}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\", + \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${SFT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${SPEC_ARGS[@]} diff --git a/setup.py b/setup.py index a1ee0475d..bf4e46f73 100644 --- a/setup.py +++ b/setup.py @@ -32,7 +32,7 @@ def get_tag(self): setup( author="vime Team", name="vime", - version="0.3.0", + version="0.3.2", packages=find_packages(include=["vime*", "vime_plugins*"]), include_package_data=True, install_requires=_fetch_requirements("requirements.txt"), diff --git a/tests/_cp_dist_helpers.py b/tests/_cp_dist_helpers.py index 1382094fe..f7b551f58 100644 --- a/tests/_cp_dist_helpers.py +++ b/tests/_cp_dist_helpers.py @@ -41,7 +41,6 @@ import sys import types - # --- Stub ``megatron.core.mpu`` (must run before cp_utils is imported) --- # # Both this module and any test file that imports it should *import this diff --git a/tests/_unit_stubs.py b/tests/_unit_stubs.py index 392ba35d8..827c615d2 100644 --- a/tests/_unit_stubs.py +++ b/tests/_unit_stubs.py @@ -35,9 +35,8 @@ def real_module_available(name: str) -> bool: """True when the real package is importable and should not be shadowed.""" - if name in sys.modules: - return True try: + # Earlier test modules may have installed a stub with no import spec. return importlib.util.find_spec(name) is not None except (ImportError, ValueError): return False @@ -104,8 +103,26 @@ def install_vllm_router_stub() -> None: return class RouterArgs: + # Stub of vllm_router.RouterArgs for CPU unit tests when the real package is absent. @classmethod - def add_cli_args(cls, parser, *args, **kwargs): # noqa: ARG003 + def add_cli_args( + cls, parser, *args, use_router_prefix=False, exclude_host_port=False, **kwargs + ): # noqa: ARG003 + prefix = "router-" if use_router_prefix else "" + dprefix = "router_" if use_router_prefix else "" + parser.add_argument( + f"--{prefix}policy", + dest=f"{dprefix}policy", + type=str, + default="cache_aware", + choices=["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash"], + ) + parser.add_argument( + f"--{prefix}request-timeout-secs", + dest=f"{dprefix}request_timeout_secs", + type=int, + default=1800, + ) return parser @classmethod @@ -170,6 +187,7 @@ def install_megatron_mpu_stub() -> MagicMock: mpu_stub.get_tensor_model_parallel_world_size.return_value = 2 mpu_stub.get_tensor_model_parallel_group.return_value = "tp_group" mpu_stub.get_pipeline_model_parallel_rank.return_value = 0 + mpu_stub.get_pipeline_model_parallel_world_size.return_value = 1 mpu_stub.get_expert_model_parallel_world_size.return_value = 1 mpu_stub.get_expert_model_parallel_group.return_value = "ep_group" @@ -209,6 +227,9 @@ def install_ray_stub() -> None: def install_vllm_cli_stubs() -> None: """Stub vLLM CLI/parser imports for ``vime.backends.vllm_utils.arguments`` when vLLM is absent.""" + # arguments.py imports RouterArgs at module load, so the router stub must be present too. + install_vllm_router_stub() + if real_module_available("vllm"): return @@ -217,6 +238,11 @@ def install_vllm_cli_stubs() -> None: utils_mod = types.ModuleType("vllm.utils") argparse_utils = types.ModuleType("vllm.utils.argparse_utils") + deep_gemm = types.ModuleType("vllm.utils.deep_gemm") + deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor = MagicMock() + deep_gemm.get_tma_aligned_size = MagicMock() + deep_gemm.is_deep_gemm_e8m0_used = MagicMock(return_value=False) + deep_gemm.per_block_cast_to_fp8 = MagicMock() import argparse @@ -225,6 +251,7 @@ class FlexibleArgumentParser(argparse.ArgumentParser): argparse_utils.FlexibleArgumentParser = FlexibleArgumentParser utils_mod.argparse_utils = argparse_utils + utils_mod.deep_gemm = deep_gemm engine_mod = types.ModuleType("vllm.engine") engine_mod.__path__ = [] @@ -237,14 +264,62 @@ def add_cli_args(cls, parser): # noqa: ARG003 arg_utils.AsyncEngineArgs = AsyncEngineArgs engine_mod.arg_utils = arg_utils + system_utils_mod = types.ModuleType("vllm.utils.system_utils") + system_utils_mod.kill_process_tree = lambda pid, include_parent=True: None # noqa: ARG005 + utils_mod.system_utils = system_utils_mod + + # vllm.entrypoints stubs (used by arguments.add_vllm_arguments and vllm_engine._vllm_server_field_names) + entrypoints_mod = types.ModuleType("vllm.entrypoints") + entrypoints_mod.__path__ = [] + openai_mod = types.ModuleType("vllm.entrypoints.openai") + openai_mod.__path__ = [] + launchers_mod = types.ModuleType("vllm.entrypoints.launchers") + launchers_mod.__path__ = [] + cli_args_mod = types.ModuleType("vllm.entrypoints.launchers.cli_args") + + import dataclasses as _dc + + @_dc.dataclass + class FrontendArgs: + @classmethod + def add_cli_args(cls, parser): # noqa: ARG003 + return parser + + cli_args_mod.FrontendArgs = FrontendArgs + cli_args_mod.make_arg_parser = lambda parser=None: parser + cli_args_mod.validate_parsed_serve_args = lambda args: args + launchers_mod.cli_args = cli_args_mod + entrypoints_mod.openai = openai_mod + entrypoints_mod.launchers = launchers_mod + vllm_mod.entrypoints = entrypoints_mod + + cli_mod = types.ModuleType("vllm.entrypoints.cli") + cli_mod.__path__ = [] + serve_mod = types.ModuleType("vllm.entrypoints.cli.serve") + + class ServeSubcommand: + pass + + serve_mod.ServeSubcommand = ServeSubcommand + cli_mod.serve = serve_mod + entrypoints_mod.cli = cli_mod + vllm_mod.engine = engine_mod vllm_mod.utils = utils_mod sys.modules["vllm"] = vllm_mod sys.modules["vllm.utils"] = utils_mod sys.modules["vllm.utils.argparse_utils"] = argparse_utils + sys.modules["vllm.utils.deep_gemm"] = deep_gemm + sys.modules["vllm.utils.system_utils"] = system_utils_mod sys.modules["vllm.engine"] = engine_mod sys.modules["vllm.engine.arg_utils"] = arg_utils + sys.modules["vllm.entrypoints"] = entrypoints_mod + sys.modules["vllm.entrypoints.openai"] = openai_mod + sys.modules["vllm.entrypoints.launchers"] = launchers_mod + sys.modules["vllm.entrypoints.launchers.cli_args"] = cli_args_mod + sys.modules["vllm.entrypoints.cli"] = cli_mod + sys.modules["vllm.entrypoints.cli.serve"] = serve_mod def install_triton_stub() -> None: diff --git a/vime/rollout/_fanout_test_helpers.py b/tests/fanout_test_helpers.py similarity index 77% rename from vime/rollout/_fanout_test_helpers.py rename to tests/fanout_test_helpers.py index 9080ff281..4ce2216ed 100644 --- a/vime/rollout/_fanout_test_helpers.py +++ b/tests/fanout_test_helpers.py @@ -1,35 +1,33 @@ -"""Test-internal compact-rollout helpers used by ``test_qwen2.5_0.5B_fanout_short.py``. +"""Compact-rollout helpers for ``test_qwen2.5_0.5B_fanout_short.py``. -The underscore prefix marks this as test infrastructure — it is not part -of the user-facing vime API and is not re-exported anywhere. It lives -in ``vime/`` only so the test can reference it by a dotted module path -(``--custom-generate-function-path`` / ``--custom-reward-post-process-path`` -resolve a string via ``importlib.import_module``, which can't handle the -dots in the e2e test's filename). +These helpers are imported by module path from the Ray job started by the E2E +test. They live on the test-only portion of ``PYTHONPATH`` because they are +test fixtures, not part of vime's public or internal runtime API. Two helpers: - ``compact_generate``: fans one input sample out to N siblings sharing the same ``rollout_id``. That's the contract the rest of the framework (per-rollout step splitter, per-rollout-mean reducer, - ``_validate_rollout_id_annotated`` validator) is built around. + ``validate_rollout_id_annotated`` validator) is built around. - ``grpo_normalize_by_group_index``: replaces the default ``_post_process_rewards`` reshape-by-shape logic with a proper ``group_index``-keyed grouping. The default at - ``vime/ray/rollout.py:_post_process_rewards`` assumes every prompt - produced exactly ``n_samples_per_prompt`` samples and reshapes by - that constant; when compact/fanout makes the per-prompt count uneven, - the reshape fails and the fallback ``view(-1, total)`` collapses - everything into ONE group, destroying per-prompt centering. - ``group_index`` (set by the data source per-prompt, preserved through - ``deepcopy``) is the right key here. + ``vime/ray/rollout.py:618`` assumes every prompt produced exactly + ``n_samples_per_prompt`` samples and reshapes by that constant; when + compact/fanout makes the per-prompt count uneven, the reshape fails + and the fallback ``view(-1, total)`` collapses everything into ONE + group, destroying per-prompt centering. ``group_index`` (set by the + data source per-prompt, preserved through ``deepcopy``) is the right + key here. """ import copy import os from collections import defaultdict + MAX_FANOUT = 3 # Each invocation appends one line. The test file reads this after train @@ -41,7 +39,7 @@ async def compact_generate(args, sample, sampling_params): """One prompt → N siblings, deterministic N = 1 + (index % MAX_FANOUT). - Strategy: call vLLM once, deepcopy N-1 times. Bounded GPU cost — + Strategy: call vllm once, deepcopy N-1 times. Bounded GPU cost — we're pinning the framework's per-rollout handling, not generation diversity. """ @@ -75,7 +73,7 @@ async def compact_generate(args, sample, sampling_params): def grpo_normalize_by_group_index(args, samples): """Drop-in ``--custom-reward-post-process-path`` for compact/fanout. - The default ``_post_process_rewards`` (``vime/ray/rollout.py``) + The default ``_post_process_rewards`` (``vime/ray/rollout.py:618``) reshapes the flat reward tensor as ``(-1, n_samples_per_prompt)`` when ``total == n_samples_per_prompt * rollout_batch_size``, falling back to ``view(-1, total)`` (= one giant group) otherwise. With diff --git a/tests/glm52_layerwise_comparator.py b/tests/glm52_layerwise_comparator.py new file mode 100644 index 000000000..015020cd3 --- /dev/null +++ b/tests/glm52_layerwise_comparator.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +"""Compare matching Megatron and vLLM decoder-layer outputs for tests.""" + +from __future__ import annotations + +import argparse +import json +import re +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +_LAYER_KEY_RE = re.compile(r"(?:^|\.)layers\.(\d+)$") +_NON_ROLLOUT_REQUEST_PREFIXES = ("HEALTH_CHECK_",) + + +@dataclass +class TrainSequence: + tokens: torch.Tensor + layers: dict[int, torch.Tensor] + source: str + + +def _load_records(path: Path) -> Iterator[dict[str, Any]]: + value = torch.load(path, map_location="cpu", weights_only=False, mmap=True) + if isinstance(value, dict): + yield value + return + if isinstance(value, list) and all(isinstance(item, dict) for item in value): + yield from value + return + raise TypeError(f"Unsupported layer dump payload in {path}: {type(value)}") + + +def _find_suffix(record: dict[str, Any], suffix: str, *, required: bool = True): + matches = [value for key, value in record.items() if key.endswith(suffix)] + if len(matches) == 1: + return matches[0] + if not matches and not required: + return None + raise KeyError(f"Expected one key ending in {suffix!r}, found {len(matches)}") + + +def _as_tensor(value: Any) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)): + tensors = [item for item in value if isinstance(item, torch.Tensor)] + if tensors: + return tensors[0] + raise TypeError(f"Expected a tensor layer output, got {type(value)}") + + +def _token_rows(value: Any, num_tokens: int, context: str) -> torch.Tensor: + tensor = _as_tensor(value) + if tensor.ndim < 2: + raise ValueError(f"{context} must have a hidden dimension, got {tensor.shape}") + rows = tensor.reshape(-1, tensor.shape[-1]) + if rows.shape[0] != num_tokens: + raise ValueError(f"{context} has {rows.shape[0]} token rows, expected {num_tokens}") + return rows + + +def _vllm_layer_token_rows(value: Any, num_tokens: int, context: str) -> torch.Tensor: + if not isinstance(value, (tuple, list)) or len(value) < 2: + raise TypeError(f"{context} must contain the vLLM layer delta and residual tensors") + delta, residual = value[:2] + if not isinstance(delta, torch.Tensor) or not isinstance(residual, torch.Tensor): + raise TypeError(f"{context} contains non-tensor layer outputs") + if delta.shape != residual.shape or delta.dtype != residual.dtype: + raise ValueError( + f"{context} delta/residual mismatch: " f"{delta.shape}/{delta.dtype} != {residual.shape}/{residual.dtype}" + ) + layer_output = (delta.float() + residual.float()).to(delta.dtype) + return _token_rows(layer_output, num_tokens, context) + + +def _layer_outputs(record: dict[str, Any], selected_layers: set[int]): + outputs = {} + for key, value in record.items(): + match = _LAYER_KEY_RE.search(key) + if match is None: + continue + layer_id = int(match.group(1)) + if layer_id in selected_layers: + outputs[layer_id] = value + return outputs + + +def load_train_sequences(dump_dir: Path, selected_layers: set[int]) -> list[TrainSequence]: + dump_files = sorted(dump_dir.glob("rank*/actor_Pass*.pt")) + if not dump_files: + raise FileNotFoundError(f"No Megatron layer dumps found under {dump_dir}") + + sequences = [] + for dump_file in dump_files: + for record in _load_records(dump_file): + tokens = record["input_ids"].reshape(-1).to(torch.int64) + cu_seqlens = record["cu_seqlens"].reshape(-1).tolist() + layer_values = record.get("layers", {}) + missing = selected_layers - set(layer_values) + if missing: + raise KeyError(f"{dump_file} is missing Megatron layers {sorted(missing)}") + layer_rows = { + layer_id: _token_rows(layer_values[layer_id], len(tokens), f"{dump_file}:layer{layer_id}") + for layer_id in selected_layers + } + for sequence_index, (start, end) in enumerate(zip(cu_seqlens, cu_seqlens[1:], strict=False)): + sequence_tokens = tokens[start:end] + if sequence_tokens.numel() == 0 or torch.count_nonzero(sequence_tokens) == 0: + continue + sequences.append( + TrainSequence( + tokens=sequence_tokens, + layers={layer_id: rows[start:end] for layer_id, rows in layer_rows.items()}, + source=f"{dump_file}:sequence{sequence_index}", + ) + ) + if not sequences: + raise RuntimeError("Megatron dumps contained no non-padding sequences") + return sequences + + +def _vllm_segments(record: dict[str, Any]): + input_ids = _find_suffix(record, ".forward_batch_info.input_ids").reshape(-1) + positions = _find_suffix(record, ".forward_batch_info.positions").reshape(-1) + rids = _find_suffix(record, ".forward_batch_info.rids") + if not rids: + return input_ids, positions, [] + + extend_seq_lens = _find_suffix(record, ".forward_batch_info.extend_seq_lens", required=False) + if extend_seq_lens is None: + counts = [1] * len(rids) + else: + counts = [int(value) for value in extend_seq_lens.reshape(-1).tolist()] + if len(counts) != len(rids) or sum(counts) != input_ids.numel(): + raise ValueError( + "vLLM request segmentation mismatch: " f"rids={len(rids)}, counts={counts}, tokens={input_ids.numel()}" + ) + + segments = [] + start = 0 + for rid, count in zip(rids, counts, strict=True): + end = start + count + rid = str(rid) + if count and not rid.startswith(_NON_ROLLOUT_REQUEST_PREFIXES): + segments.append((rid, slice(start, end))) + start = end + return input_ids.to(torch.int64), positions.to(torch.int64), segments + + +def _vllm_dump_files(dump_dir: Path) -> list[Path]: + process_dirs = sorted(path for path in dump_dir.iterdir() if path.is_dir()) + files = [] + for process_dir in process_dirs: + files.extend(sorted(process_dir.glob("Chunk*.pt"))) + files.extend(sorted(process_dir.glob("Pass*.pt"))) + if not files: + raise FileNotFoundError(f"No vLLM layer dumps found under {dump_dir}") + return files + + +def map_requests_to_train_sequences(dump_files: list[Path], train_sequences: list[TrainSequence]) -> dict[str, int]: + observations: dict[str, dict[int, int]] = {} + for dump_file in dump_files: + for record in _load_records(dump_file): + input_ids, positions, segments = _vllm_segments(record) + for rid, token_slice in segments: + request_observations = observations.setdefault(rid, {}) + for position, token_id in zip( + positions[token_slice].tolist(), + input_ids[token_slice].tolist(), + strict=True, + ): + previous = request_observations.setdefault(position, token_id) + if previous != token_id: + raise ValueError( + f"vLLM request {rid} changed token at position {position}: " f"{previous} != {token_id}" + ) + + mapping = {} + for rid, request_observations in observations.items(): + candidates = [] + for sequence_id, sequence in enumerate(train_sequences): + if all( + 0 <= position < sequence.tokens.numel() and int(sequence.tokens[position]) == token_id + for position, token_id in request_observations.items() + ): + candidates.append(sequence_id) + if not candidates: + raise RuntimeError(f"Could not map vLLM request {rid} to any Megatron token sequence") + mapping[rid] = candidates[0] + if not mapping: + raise RuntimeError("vLLM dumps contained no request observations") + return mapping + + +def compare_layer_outputs( + dump_files: list[Path], + train_sequences: list[TrainSequence], + request_mapping: dict[str, int], + selected_layers: set[int], +): + stats = {layer_id: {"max_abs": 0.0, "sum_abs": 0.0, "numel": 0, "tokens": 0} for layer_id in selected_layers} + compared: set[tuple[str, int, int]] = set() + + for dump_file in dump_files: + for record in _load_records(dump_file): + input_ids, positions, segments = _vllm_segments(record) + if not segments: + continue + rollout_layers = _layer_outputs(record, selected_layers) + missing = selected_layers - set(rollout_layers) + if missing: + raise KeyError(f"{dump_file} is missing vLLM layers {sorted(missing)}") + rollout_rows = { + layer_id: _vllm_layer_token_rows( + value, + input_ids.numel(), + f"{dump_file}:layer{layer_id}", + ) + for layer_id, value in rollout_layers.items() + } + + for rid, token_slice in segments: + train_sequence = train_sequences[request_mapping[rid]] + segment_positions = positions[token_slice] + for layer_id in selected_layers: + keep = torch.tensor( + [ + # A causal LM never consumes the hidden state at the + # final input position to score a token in this + # sequence. vLLM may still execute that terminal + # token after sampling it, whereas Megatron's + # log-prob forward stops at score-producing + # positions. The terminal state therefore has no + # corresponding training row to compare. + int(position) < train_sequence.tokens.numel() - 1 + and (rid, int(position), layer_id) not in compared + for position in segment_positions.tolist() + ], + dtype=torch.bool, + ) + if not torch.any(keep): + continue + kept_positions = segment_positions[keep] + if torch.any(kept_positions < 0) or torch.any(kept_positions >= train_sequence.tokens.numel()): + raise IndexError(f"vLLM request {rid} contains positions outside its " "Megatron sequence") + rollout_value = rollout_rows[layer_id][token_slice][keep] + train_value = train_sequence.layers[layer_id][kept_positions] + difference = (rollout_value.float() - train_value.float()).abs() + layer_stats = stats[layer_id] + layer_stats["max_abs"] = max(layer_stats["max_abs"], float(difference.max().item())) + layer_stats["sum_abs"] += float(difference.sum().item()) + layer_stats["numel"] += difference.numel() + layer_stats["tokens"] += kept_positions.numel() + for position in kept_positions.tolist(): + compared.add((rid, int(position), layer_id)) + + for layer_stats in stats.values(): + layer_stats["mean_abs"] = ( + layer_stats.pop("sum_abs") / layer_stats["numel"] if layer_stats["numel"] else float("nan") + ) + return stats + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--megatron-dir", type=Path, required=True) + parser.add_argument("--vllm-dir", type=Path, required=True) + parser.add_argument("--layers", type=int, nargs="+", required=True) + parser.add_argument("--max-hidden-diff", type=float, default=1e-7) + parser.add_argument("--num-threads", type=int, default=1) + parser.add_argument("--output-json", type=Path) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.num_threads <= 0: + raise ValueError(f"--num-threads must be positive, got {args.num_threads}") + # Records contain only a few token rows. Launching the host-wide PyTorch + # thread pool for every small reduction costs far more than the arithmetic. + torch.set_num_threads(args.num_threads) + selected_layers = set(args.layers) + train_sequences = load_train_sequences(args.megatron_dir, selected_layers) + dump_files = _vllm_dump_files(args.vllm_dir) + request_mapping = map_requests_to_train_sequences(dump_files, train_sequences) + stats = compare_layer_outputs( + dump_files, + train_sequences, + request_mapping, + selected_layers, + ) + + failed = False + for layer_id in sorted(stats): + layer_stats = stats[layer_id] + print( + f"layer={layer_id} tokens={layer_stats['tokens']} " + f"max_abs={layer_stats['max_abs']:.12g} " + f"mean_abs={layer_stats['mean_abs']:.12g}" + ) + if not layer_stats["tokens"] or layer_stats["max_abs"] > args.max_hidden_diff: + failed = True + + result = { + "max_hidden_diff": args.max_hidden_diff, + "request_mapping": request_mapping, + "layers": stats, + } + if args.output_json is not None: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(result, indent=2) + "\n") + if failed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/utils/test_trace_utils.py b/tests/observability/test_trace_utils.py similarity index 52% rename from tests/utils/test_trace_utils.py rename to tests/observability/test_trace_utils.py index 162c36924..4bcba095b 100644 --- a/tests/utils/test_trace_utils.py +++ b/tests/observability/test_trace_utils.py @@ -5,9 +5,11 @@ import pytest import torch -from vime.utils.trace_utils import trace_span +from vime.observability.trace_utils import TRACE_CHILDREN_KEY, build_vllm_meta_trace_attrs, trace_span from vime.utils.types import Sample +NUM_GPUS = 0 + def _load_trace_timeline_viewer_module(): module_path = Path(__file__).resolve().parents[2] / "tools" / "trace_timeline_viewer.py" @@ -22,6 +24,37 @@ def _load_trace_timeline_viewer_module(): return module +@pytest.mark.unit +def test_build_vllm_meta_trace_attrs_keeps_standard_and_pd_fields(): + meta = { + "prompt_tokens": 12, + "completion_tokens": 7, + "cached_tokens": 3, + "pd_prefill_forward_duration": 0.125, + "pd_decode_transfer_duration": 0.05, + "finish_reason": {"type": "stop"}, + "unused_field": "ignored", + } + + attrs = build_vllm_meta_trace_attrs(meta) + trace_children = attrs.pop(TRACE_CHILDREN_KEY) + + assert attrs == { + "prompt_tokens": 12, + "completion_tokens": 7, + "cached_tokens": 3, + "finish_reason": "stop", + } + assert trace_children[0]["name"] == "vllm_pd_prefill" + assert trace_children[0]["children"][0]["attrs"] == { + "pd_prefill_forward_duration": 0.125, + } + assert trace_children[1]["name"] == "vllm_pd_decode" + assert trace_children[1]["children"][0]["attrs"] == { + "pd_decode_transfer_duration": 0.05, + } + + @pytest.mark.unit def test_trace_timeline_viewer_omits_virtual_pd_lanes_without_pd_attrs(tmp_path: Path): viewer = _load_trace_timeline_viewer_module() @@ -29,12 +62,14 @@ def test_trace_timeline_viewer_omits_virtual_pd_lanes_without_pd_attrs(tmp_path: with trace_span(sample, "vllm_generate", attrs={"max_new_tokens": 8}) as span: span.update( - { - "prompt_tokens": 4, - "completion_tokens": 2, - "cached_tokens": 1, - "finish_reason": "stop", - } + build_vllm_meta_trace_attrs( + { + "prompt_tokens": 4, + "completion_tokens": 2, + "cached_tokens": 1, + "finish_reason": {"type": "stop"}, + } + ) ) pt_path = tmp_path / "rollout.pt" @@ -58,3 +93,7 @@ def test_trace_timeline_viewer_omits_virtual_pd_lanes_without_pd_attrs(tmp_path: } assert "[P]" not in item["name"] assert "[D]" not in item["name"] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/plugin_contracts/test_plugin_path_loading_contracts.py b/tests/plugin_contracts/test_plugin_path_loading_contracts.py index 9a1fb2d9a..a55100f90 100644 --- a/tests/plugin_contracts/test_plugin_path_loading_contracts.py +++ b/tests/plugin_contracts/test_plugin_path_loading_contracts.py @@ -34,7 +34,11 @@ from vime.rollout.base_types import RolloutFnEvalOutput, call_rollout_fn from vime.rollout.data_source import RolloutDataSourceWithBuffer -from vime.rollout.filter_hub.base_types import DynamicFilterOutput, call_dynamic_filter +from vime.rollout.filter_hub.base_types import ( + DynamicFilterOutput, + call_dynamic_filter, + should_drop_dynamic_filter_output, +) from vime.rollout.rm_hub import async_rm, batched_async_rm from vime.rollout.vllm_rollout import generate_rollout as default_generate_rollout from vime.utils.misc import load_function @@ -198,6 +202,25 @@ def check_dynamic_filter_path(path: str) -> None: assert isinstance(output, DynamicFilterOutput) +def test_dynamic_filter_can_fall_back_to_zero_std_groups_at_target_capacity(): + fn = load_function("vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std_with_fallback") + output = call_dynamic_filter(fn, make_args(), [make_sample(0, reward=1.0), make_sample(1, reward=1.0)]) + + assert not bool(output.keep) + assert output.keep_when_insufficient is True + assert should_drop_dynamic_filter_output(output, remaining_batch_size=3, target_data_size=2) is True + assert should_drop_dynamic_filter_output(output, remaining_batch_size=2, target_data_size=2) is False + + +def test_strict_dynamic_filter_still_drops_at_target_capacity(): + fn = load_function("vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std") + output = call_dynamic_filter(fn, make_args(), [make_sample(0, reward=0.0), make_sample(1, reward=0.0)]) + + assert not bool(output.keep) + assert output.keep_when_insufficient is False + assert should_drop_dynamic_filter_output(output, remaining_batch_size=2, target_data_size=2) is True + + def check_buffer_filter_default() -> None: fn = load_function("vime.rollout.data_source.pop_first") assert tuple(inspect.signature(fn).parameters)[:4] == ("args", "rollout_id", "buffer", "num_samples") diff --git a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py index 3c066996f..56310cabf 100644 --- a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py +++ b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py @@ -133,7 +133,7 @@ def invoke_rollout_data_postprocess(fn): "custom_rollout_log", "CUSTOM_ROLLOUT_LOG_FUNCTION_PATH", "plugin_contracts.test_plugin_runtime_hook_contracts.reference_custom_rollout_log", - "vime/ray/rollout.py", + "vime/observability/rollout_metrics.py", "custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time)", ("rollout_id", "args", "samples", "rollout_extra_metrics", "rollout_time"), invoke_custom_rollout_log, @@ -142,7 +142,7 @@ def invoke_rollout_data_postprocess(fn): "custom_eval_rollout_log", "CUSTOM_EVAL_ROLLOUT_LOG_FUNCTION_PATH", "plugin_contracts.test_plugin_runtime_hook_contracts.reference_custom_eval_rollout_log", - "vime/ray/rollout.py", + "vime/observability/rollout_metrics.py", "custom_log_func(rollout_id, args, data, extra_metrics)", ("rollout_id", "args", "data", "extra_metrics"), invoke_custom_eval_rollout_log, diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py new file mode 100644 index 000000000..7826695fb --- /dev/null +++ b/tests/test_accelerator.py @@ -0,0 +1,203 @@ +from types import SimpleNamespace + +import pytest + +from vime.utils import accelerator + +NUM_GPUS = 0 + + +class FakeAccelerator(accelerator.Accelerator): + name = "fake" + device_type = "fake" + communication_backend_name = "fake" + + def is_available(self): + return True + + def device(self, index=None): + return accelerator.torch.device("cpu") + + def device_name(self, index=None): + return "cpu" + + def set_device(self, index): + return None + + def current_device(self): + return "cpu" + + def device_count(self): + return 0 + + def synchronize(self, device=None): + return None + + def current_stream(self, device=None): + return None + + def empty_cache(self): + return None + + def mem_get_info(self, device=None): + return 0, 0 + + def memory_allocated(self, device=None): + return 0 + + def memory_reserved(self, device=None): + return 0 + + +@pytest.fixture(autouse=True) +def reset_accelerator_selection(monkeypatch): + registry = accelerator._REGISTRY.copy() + selected = accelerator._ACCELERATOR + patch_imported = accelerator._MUSA_PATCH_IMPORTED + bootstrap_checked = accelerator._MUSA_BOOTSTRAP_CHECKED + for name in ("VIME_ACCELERATOR", "MUSA_VISIBLE_DEVICES", "MUSA_PATCH_PATH", "CUDA_VISIBLE_DEVICES"): + monkeypatch.delenv(name, raising=False) + accelerator._REGISTRY.clear() + accelerator.reset_accelerator() + accelerator._MUSA_PATCH_IMPORTED = False + accelerator._MUSA_BOOTSTRAP_CHECKED = False + yield + accelerator._REGISTRY.clear() + accelerator._REGISTRY.update(registry) + accelerator._ACCELERATOR = selected + accelerator._MUSA_PATCH_IMPORTED = patch_imported + accelerator._MUSA_BOOTSTRAP_CHECKED = bootstrap_checked + + +@pytest.mark.unit +def test_cuda_selection_does_not_bootstrap_musa(monkeypatch): + monkeypatch.setenv("VIME_ACCELERATOR", "cuda") + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + monkeypatch.setattr(accelerator, "_cuda_available", lambda: True) + monkeypatch.setattr(accelerator.CUDAAccelerator, "is_available", lambda self: True) + monkeypatch.setattr( + accelerator, + "_import_musa_patch", + lambda: pytest.fail("CUDA selection must not import musa_patch"), + ) + + assert accelerator.get_accelerator().name == "cuda" + assert accelerator.process_group_backend() == "nccl" + assert accelerator.visible_devices_env_key() == "CUDA_VISIBLE_DEVICES" + + +@pytest.mark.unit +def test_selected_musa_bootstraps_patch_once(monkeypatch): + imports = [] + fake_musa = SimpleNamespace(is_available=lambda: True) + + def import_musa_patch(): + imports.append("musa_patch") + monkeypatch.setattr(accelerator.torch, "musa", fake_musa, raising=False) + return True + + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + monkeypatch.setattr(accelerator, "_import_musa_patch", import_musa_patch) + + assert imports == [] + assert accelerator.initialize_accelerator().name == "musa" + assert accelerator.initialize_accelerator().name == "musa" + assert imports == ["musa_patch"] + + +@pytest.mark.unit +def test_cpu_only_initialization_does_not_require_an_accelerator(monkeypatch): + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + monkeypatch.setattr(accelerator, "_cuda_available", lambda: False) + + assert accelerator.initialize_accelerator() is None + + +@pytest.mark.unit +def test_musa_backend_maps_devices_and_process_groups(monkeypatch): + monkeypatch.setattr(accelerator.MUSAAccelerator, "is_available", lambda self: True) + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "2,5") + accelerator.set_accelerator(accelerator.MUSAAccelerator()) + + assert accelerator.visible_devices_env_key() == "MUSA_VISIBLE_DEVICES" + assert accelerator.resolve_visible_device_id("5") == 1 + assert accelerator.process_group_backend() == "mccl" + assert accelerator.weight_update_backend() == "cpu:gloo,musa:mccl" + assert accelerator.process_group_backend("gloo") == "gloo" + + +@pytest.mark.unit +def test_cuda_visible_device_mapping(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,6") + accelerator.set_accelerator(FakeAccelerator()) + + assert accelerator.resolve_visible_device_id(4) == 0 + assert accelerator.resolve_visible_device_id(1) == 1 + with pytest.raises(RuntimeError, match="CUDA_VISIBLE_DEVICES=4,6"): + accelerator.resolve_visible_device_id(7) + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-aaaa,GPU-bbbb") + assert accelerator.resolve_visible_device_id("GPU-bbbb") == 1 + + +@pytest.mark.unit +def test_registered_backend_can_be_selected(monkeypatch): + class RegisteredAccelerator(FakeAccelerator): + name = "registered" + + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + monkeypatch.setattr(accelerator, "_cuda_available", lambda: False) + accelerator.register_accelerator("registered", RegisteredAccelerator, lambda: True, priority=300) + + assert accelerator.get_accelerator().name == "registered" + + +@pytest.mark.unit +def test_cuda_backend_uses_torch_cuda_namespace(monkeypatch): + monkeypatch.setattr(accelerator.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(accelerator.torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(accelerator.torch.cuda, "current_device", lambda: 1) + monkeypatch.setattr(accelerator.torch.cuda, "memory_allocated", lambda device=None: 123) + backend = accelerator.CUDAAccelerator() + + assert backend.is_available() + assert backend.device_name() == "cuda:1" + assert backend.memory_allocated() == 123 + + +@pytest.mark.unit +def test_routing_replay_uses_selected_backend_current_device(monkeypatch): + from vime.utils import routing_replay + + transfers = [] + + class FakeTopIndices: + def is_pinned(self): + return False + + def to(self, device, *, dtype, non_blocking): + transfers.append((device, dtype, non_blocking)) + return self + + accelerator.set_accelerator(FakeAccelerator()) + monkeypatch.setattr(routing_replay.RoutingReplay, "all_routing_replays", []) + replay = routing_replay.RoutingReplay() + replay.top_indices_list.append(FakeTopIndices()) + + replay.pop_forward() + replay.pop_backward() + assert transfers == [ + ("cpu", accelerator.torch.int32, False), + ("cpu", accelerator.torch.int32, False), + ] + + +@pytest.mark.unit +def test_musa_availability_handles_missing_torch_namespace(monkeypatch): + monkeypatch.delattr(accelerator.torch, "musa", raising=False) + + assert accelerator.is_musa_available() is False + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_advantage_whiten_cp.py b/tests/test_advantage_whiten_cp.py new file mode 100644 index 000000000..e0589b51d --- /dev/null +++ b/tests/test_advantage_whiten_cp.py @@ -0,0 +1,192 @@ +"""Advantage-whitening CP-invariance check on CPU. + +``compute_advantages_and_returns`` whitens advantages with +``distributed_masked_whiten``, which all-reduces ``(sum, sum_sq, mask_sum)`` +over the process group it is handed. Under context parallelism each rank only +holds its zigzag slice of every sequence, so those statistics are only the +*global* statistics if the group spans the CP dimension as well as DP. + +If the CP-excluding group is used instead, every CP rank normalizes with the +mean/variance of its own slice: the two halves of one sequence come out with +different affine transforms, and none of them matches the correct whitening. + +The contract pinned here: for a fixed set of samples, the whitened advantage of +a given sample is the same number for every ``(dp_size, cp_size)`` factorization +of the world — and in particular the same as the single-rank baseline. + +Two of the cases use a prompt-heavy sequence whose response tokens all land on +one CP rank, so some ranks contribute an entirely empty local mask. Those ranks +must still take part in the all-reduce; a rank-dependent "skip whitening when I +hold nothing" shortcut desyncs the collective and hangs. The spawn join below +is bounded so that regression surfaces as a failure rather than a stuck job. +""" + +from __future__ import annotations + +import json +import os +import time + +# Megatron stub must land in sys.modules before anything imports +# vime.backends.megatron_utils. pytest's prepend importmode puts ``tests/`` on +# sys.path, so the bare-name import works without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +from _cp_dist_helpers import free_port, stub_megatron_in_worker + + +NUM_GPUS = 0 + +# (total_length, response_length) per sample, plus its rollout reward. Eight +# samples so dp_size in {1, 2, 4} divides evenly. Sample 3 is deliberately +# prompt-heavy: at cp_size >= 2 its response sits entirely inside one rank's +# chunk, leaving the other ranks with an empty mask for it. +SEQS = [(64, 48), (100, 90), (40, 12), (256, 16), (72, 30), (90, 84), (50, 20), (110, 96)] +REWARDS = [1.5, -0.5, -1.0, 2.0, 0.25, -1.75, 0.75, -0.25] + +WHITEN_CASES = [(1, 1), (2, 1), (1, 2), (2, 2), (1, 4), (4, 1)] + + +class _Args: + advantage_estimator = "grpo" + normalize_advantages = True + kl_coef = 0.0 + use_kl_loss = False + use_rollout_logprobs = False + use_opd = False + custom_advantage_function_path = None + kl_loss_type = "low_var_kl" + gamma = 1.0 + lambd = 1.0 + + +def _whiten_worker(rank, world_size, cp_size, dp_size, master_port, result_dir): + """One spawned rank: whiten its DP shard's CP slice, dump the results.""" + import torch + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + _dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + try: + from megatron.core import mpu + + # No TP/PP here, so DP-with-CP is the whole world. The DP-only group is + # the set of ranks sharing this rank's cp_rank -- exactly the group that + # Megatron's ``get_data_parallel_group(with_context_parallel=False)`` + # returns. Every rank must build every subgroup, in the same order. + dp_cp_group = _dist.new_group(ranks=list(range(world_size))) + dp_only_groups = [ + _dist.new_group(ranks=[r for r in range(world_size) if r % cp_size == c]) for c in range(cp_size) + ] + + mpu.is_pipeline_last_stage = lambda: True + mpu.get_data_parallel_group = lambda with_context_parallel=False, **kw: ( + dp_cp_group if with_context_parallel else dp_only_groups[cp_rank] + ) + + from vime.backends.megatron_utils.cp_utils import get_logits_and_tokens_offset_with_cp + from vime.backends.megatron_utils.loss import compute_advantages_and_returns + + def cp_slice(x, total_len, response_len): + """Keep only the response positions this CP rank owns.""" + if cp_size == 1: + return x + prompt_len = total_len - response_len + _, _, _, offsets = get_logits_and_tokens_offset_with_cp(total_len, response_len) + parts = [] + for start, end in offsets: + lo, hi = max(0, start - prompt_len), max(0, end - prompt_len) + if hi > lo: + parts.append(x[lo:hi]) + return torch.cat(parts) if parts else x[:0] + + # Round-robin DP shard, mirroring how samples spread over DP ranks. + my_samples = [i for i in range(len(SEQS)) if i % dp_size == dp_rank] + + rollout_data = { + "log_probs": [cp_slice(torch.zeros(SEQS[i][1]), *SEQS[i]) for i in my_samples], + "loss_masks": [torch.ones(SEQS[i][1]) for i in my_samples], + "rewards": [REWARDS[i] for i in my_samples], + "response_lengths": [SEQS[i][1] for i in my_samples], + "total_lengths": [SEQS[i][0] for i in my_samples], + } + + compute_advantages_and_returns(_Args(), rollout_data) + + # grpo gives every token of a sample the same advantage, and whitening is + # affine, so one value per sample fully describes the result. Ranks that + # own no tokens of a sample simply report nothing for it. + out = { + str(i): round(adv[0].item(), 6) + for i, adv in zip(my_samples, rollout_data["advantages"], strict=True) + if adv.numel() > 0 + } + with open(os.path.join(result_dir, f"rank{rank}.json"), "w") as f: + json.dump(out, f) + finally: + _dist.destroy_process_group() + + +def _run_case(dp_size, cp_size, tmp_path): + """Spawn the world, then merge every rank's per-sample whitened values.""" + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_dir = tmp_path / f"dp{dp_size}_cp{cp_size}" + result_dir.mkdir(parents=True, exist_ok=True) + + ctx = mp.spawn( + _whiten_worker, + args=(world_size, cp_size, dp_size, free_port(), str(result_dir)), + nprocs=world_size, + join=False, + ) + deadline = time.time() + 180 + while not ctx.join(timeout=5): + if time.time() > deadline: + for p in ctx.processes: + if p.is_alive(): + p.terminate() + pytest.fail( + f"dp={dp_size} cp={cp_size} workers did not finish in 180s; " + "a rank-dependent branch around the whitening all_reduce desyncs the collective" + ) + + merged: dict[str, float] = {} + for rank in range(world_size): + with open(result_dir / f"rank{rank}.json") as f: + for sample_idx, value in json.load(f).items(): + if sample_idx in merged: + # Every rank holding part of a sample must agree on it -- + # this is the assertion that fails on the CP-excluding group. + assert merged[sample_idx] == pytest.approx(value, abs=1e-5), ( + f"dp={dp_size} cp={cp_size}: CP ranks disagree on sample {sample_idx}: " + f"{merged[sample_idx]} vs {value}" + ) + merged[sample_idx] = value + return merged + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_size,cp_size", WHITEN_CASES) +def test_whitened_advantages_are_cp_invariant(dp_size, cp_size, tmp_path): + baseline = _run_case(1, 1, tmp_path) + assert len(baseline) == len(SEQS), "baseline should cover every sample" + + got = _run_case(dp_size, cp_size, tmp_path) + + assert sorted(got) == sorted(baseline) + for sample_idx, expected in baseline.items(): + assert got[sample_idx] == pytest.approx(expected, abs=1e-5), ( + f"dp={dp_size} cp={cp_size}: sample {sample_idx} whitened to {got[sample_idx]}, " + f"single-rank baseline is {expected}" + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/unit/__init__.py b/tests/test_agent/__init__.py similarity index 100% rename from tests/unit/__init__.py rename to tests/test_agent/__init__.py diff --git a/tests/test_agent/_dump_helpers.py b/tests/test_agent/_dump_helpers.py new file mode 100644 index 000000000..6dec5554d --- /dev/null +++ b/tests/test_agent/_dump_helpers.py @@ -0,0 +1,85 @@ +"""Tree dumper for the trajectory_manager branching test. + +Renders ``TrajectoryManager._trees[sid]`` as ASCII text so a human can read the +routing tree next to the linearized Samples. Adapted to the refactored +``MessageNode`` API (``.message`` / ``.turn`` / ``.turn_index``); see +``vime/agent/trajectory.py``. + +Pulled out of the historic ``tests/test_coding_agent/_dump_helpers.py`` (which +also carried JSON dumpers + an aiohttp debug middleware) down to just the one +helper the branching test imports. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + + +def dump_tree_txt(manager, sid: str, *, max_text_chars: int = 80) -> str: + """Render ``manager._trees[sid]`` as ASCII text. + + Returns ``""`` when the sid has no tree (already drained + or never opened). Otherwise a header line plus one indented row per non-root + node, listing role / message text and, for generated assistant leaves, the + turn snapshot fields (turn index, prompt/response id counts, finish reason). + """ + root = manager._trees.get(sid) + if root is None: + return f"" + non_root_count = sum(1 for _ in _iter_non_root(root)) + leaf_count = sum(1 for leaf in root.leaves() if not leaf.is_root) + lines: list[str] = [ + f"session={sid} turns={manager.turn_count(sid)} leaves={leaf_count} nodes={non_root_count}", + "root", + ] + _render_subtree(root, lines, depth=0, max_text_chars=max_text_chars) + return "\n".join(lines) + + +def _iter_non_root(root) -> Iterator: + stack: list = list(root.children) + while stack: + n = stack.pop() + yield n + stack.extend(n.children) + + +def _message_text(message: dict[str, Any] | None) -> str: + """Human-readable text for one node's chat message (``None`` for root / + empty-response leaves).""" + if not message: + return "" + content = message.get("content") + parts: list[str] = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append("reason:" + reasoning) + if isinstance(content, str): + parts.append(content) + elif isinstance(content, list): + for b in content: + if isinstance(b, dict) and b.get("type") == "text": + parts.append(str(b.get("text", ""))) + elif content is not None: + parts.append(str(content)) + return " ".join(p for p in parts if p) + + +def _render_subtree(node, lines: list[str], *, depth: int, max_text_chars: int) -> None: + indent = " " * depth + for child in node.children: + text = _message_text(child.message) + if len(text) > max_text_chars: + text = text[:max_text_chars] + "..." + text = text.replace("'", "\\'") + fields = [f"[{child.role}]"] + if child.role == "assistant" and child.turn is not None: + fields.append(f"turn={child.turn_index}") + fields.append(f"prompt_ids=({len(child.turn.prompt_ids)})") + fields.append(f"response_ids=({len(child.turn.output_ids)})") + fields.append(f"finish={child.turn.finish_reason}") + fields.append(f"has_logprobs={bool(child.turn.output_log_probs)}") + fields.append(f"text='{text}'") + lines.append(f"{indent}└── " + " ".join(fields)) + _render_subtree(child, lines, depth=depth + 1, max_text_chars=max_text_chars) diff --git a/tests/test_agent/_fakes.py b/tests/test_agent/_fakes.py new file mode 100644 index 000000000..e342ec08f --- /dev/null +++ b/tests/test_agent/_fakes.py @@ -0,0 +1,316 @@ +"""Shared CPU-only fakes for the agent test suite. + +These stand in for the four real external boundaries of an agent rollout so the +whole pipeline (generate -> sandbox -> harness -> adapter HTTP -> vllm -> +trajectory) can run deterministically on CPU, no GPU / E2B / vllm / +checkpoint required. The code under test stays real; only these edges are faked: + + * :class:`FakeTokenizer` -- word-level chat-template render + decode that + round-trips (``decode(encode(t)) == t``), so a + scripted model reply survives the + encode->generate->decode->parse round trip. + * :class:`ScriptedTokenizer` -- pre-baked prompt-id queue + id->text decode, + for adapter unit tests that assert exact ids. + * :class:`FakeVLLMServer` -- a real aiohttp ``/inference/v1/generate`` upstream returning + scripted ``output_token_logprobs`` per turn + (exercises the real HTTP path in + ``common.call_vllm_generate``). + * :func:`fake_call_vllm_generate` -- a drop-in for + ``common.call_vllm_generate`` (monkeypatch) + that skips HTTP and yields scripted TurnRecords. + * :class:`FakeSandbox` -- an in-memory :class:`vime.agent.sandbox.Sandbox` + that records every ``exec`` and drives the + detached-launch/poll handshake via ``on_launch``. +""" + +from __future__ import annotations + +import re +from collections.abc import Awaitable, Callable + +from aiohttp import web + +from vime.agent.adapters.common import TurnRecord + +# --------------------------------------------------------------------------- +# Tokenizers +# --------------------------------------------------------------------------- + +# Fixed ids for the structural chat-template markers (kept out of the dynamic +# word band so decode can drop them as "special"). +_ROLE_BEGIN = {"system": 1, "user": 2, "assistant": 3, "tool": 4} +_ROLE_END = 5 +_GEN = 3 # add_generation_prompt marker == assistant-begin (mirrors a real template) +_SPECIAL_IDS = set(_ROLE_BEGIN.values()) | {_ROLE_END, _GEN} +_WORD_BASE = 100 # dynamic word ids start here, never colliding with specials + + +class FakeTokenizer: + """Deterministic word-level tokenizer with a round-tripping decode. + + Each distinct whitespace-delimited word gets a stable id (>= ``_WORD_BASE``) + on first sight, so ``decode(encode(text)) == text`` for the content words. + ``apply_chat_template`` frames each message as ``[ROLE_BEGIN, *words, END]`` + and appends the generation marker, the same shape a real template emits -- + enough for the manager to see clean prefix extensions when an assistant turn + is replayed verbatim on the next request. + """ + + def __init__(self, outputs: dict[tuple[int, ...], str] | None = None) -> None: + self._vocab: dict[str, int] = {} + self._inv: dict[int, str] = {} + # Explicit output-id -> text map for decode, decoupled from encode: lets a + # test script the model server to return fixed ids and still control the + # decoded reply text (mirrors the historic ToyTokenizer.outputs pattern). + self._outputs = dict(outputs or {}) + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def _id(self, word: str) -> int: + if word not in self._vocab: + tid = _WORD_BASE + len(self._vocab) + self._vocab[word] = tid + self._inv[tid] = word + return self._vocab[word] + + def encode(self, text: str) -> list[int]: + return [self._id(w) for w in text.split()] if text else [] + + def decode(self, ids, skip_special_tokens: bool = False) -> str: + ids = list(ids) + if tuple(ids) in self._outputs: + return self._outputs[tuple(ids)] + return " ".join(self._inv[i] for i in ids if i in self._inv and i not in _SPECIAL_IDS) + + @staticmethod + def _content_text(message: dict) -> str: + c = message.get("content") + parts: list[str] = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append(reasoning) + if isinstance(c, str): + parts.append(c) + elif isinstance(c, list): + for b in c: + if isinstance(b, dict) and b.get("type") == "text": + parts.append(str(b.get("text", ""))) + for call in message.get("tool_calls") or []: + fn = call.get("function") or {} + parts.append(f"toolcall:{fn.get('name', '')}") + return " ".join(p for p in parts if p) + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + out: list[int] = [] + for m in messages: + role = m.get("role", "user") + out.append(_ROLE_BEGIN.get(role, _ROLE_BEGIN["user"])) + out.extend(self.encode(self._content_text(m))) + out.append(_ROLE_END) + if add_generation_prompt: + out.append(_GEN) + return out + + +class ScriptedTokenizer: + """Pre-baked prompt-id queue + id->text decode, for adapter unit tests that + assert exact token sequences. ``apply_chat_template`` ignores the messages + and pops the next scripted prompt; ``decode`` maps an output-id tuple to its + scripted text.""" + + def __init__(self, prompts: list[list[int]], outputs: dict[tuple[int, ...], str]) -> None: + self.prompts = [list(p) for p in prompts] + self.outputs = dict(outputs) + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + assert self.prompts, "unexpected chat-template render (prompt queue exhausted)" + return self.prompts.pop(0) + + def decode(self, ids, skip_special_tokens: bool = False) -> str: + return self.outputs.get(tuple(ids), "") + + +# --------------------------------------------------------------------------- +# vllm /inference/v1/generate fakes +# --------------------------------------------------------------------------- + + +class FakeVLLMServer: + """Real aiohttp ``/inference/v1/generate`` upstream returning scripted turns. + + Each turn is a list of ``(logprob, token_id)`` pairs that the server emits as + a vLLM ``choices[0]`` payload (``token_ids`` + ``logprobs.content[i].logprob`` + + a string ``finish_reason``), the shape + ``common._tokens_and_logprobs_from_choice`` parses. Records every request body + + the ``x-session-id`` routing header so tests can assert the adapter posted + the right ``token_ids`` / sampling params. Use as an async context manager; + ``.url`` is the base url to hand the adapter. + """ + + def __init__(self, turns: list[list[tuple[float, int]]], *, finish_reason: str = "stop") -> None: + self.turns = [list(t) for t in turns] + self.finish_reason = finish_reason + self.requests: list[dict] = [] + self.routing_keys: list[str | None] = [] + self._server: web.Application | None = None + self._runner = None + + async def _handle(self, request: web.Request) -> web.Response: + self.routing_keys.append(request.headers.get("x-session-id")) + self.requests.append(await request.json()) + assert self.turns, "unexpected /inference/v1/generate call (turn script exhausted)" + pairs = self.turns.pop(0) + token_ids = [tid for _lp, tid in pairs] + content = [{"logprob": lp} for lp, _tid in pairs] + return web.json_response( + { + "choices": [ + { + "token_ids": token_ids, + "logprobs": {"content": content}, + "finish_reason": self.finish_reason, + } + ] + } + ) + + async def __aenter__(self) -> FakeVLLMServer: + from aiohttp.test_utils import TestServer + + app = web.Application() + app.router.add_post("/inference/v1/generate", self._handle) + self._server = TestServer(app) + await self._server.start_server() + self.url = str(self._server.make_url("")).rstrip("/") + return self + + async def __aexit__(self, *exc) -> None: + if self._server is not None: + await self._server.close() + + +def fake_call_vllm_generate(scripted: list[tuple[str, str, list[float] | None]], tokenizer: FakeTokenizer): + """Build a drop-in for ``common.call_vllm_generate`` (for monkeypatch). + + ``scripted`` is a list of ``(response_text, finish_reason, logprobs)`` consumed + one per turn. The response text is encoded with ``tokenizer`` so the adapter's + ``decode(output_ids)`` round-trips it back, and the real + ``parse_model_output`` runs on it. The returned coroutine matches the real + signature ``(prompt_ids, session, body, *, adapter, session_id=None)``. + """ + queue = list(scripted) + + async def _fake(prompt_ids, session, body, *, adapter, session_id=None) -> TurnRecord: + assert queue, "unexpected vllm /inference/v1/generate call (response script exhausted)" + text, finish, logprobs = queue.pop(0) + output_ids = tokenizer.encode(text) + lp = list(logprobs) if logprobs is not None else [0.0] * len(output_ids) + assert len(lp) == len(output_ids), "scripted logprobs length must match encoded response" + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=output_ids, + finish_reason=finish, + output_log_probs=lp, + ) + + return _fake + + +# --------------------------------------------------------------------------- +# Sandbox +# --------------------------------------------------------------------------- + +_POLL_RE = re.compile(r"test -f (\S+) && cat \1") + + +class FakeSandbox: + """In-memory :class:`vime.agent.sandbox.Sandbox` for CPU tests. + + Records every ``exec`` (so harness tests can assert the right commands were + issued) and keeps an in-memory file store for ``write_file`` / ``read_file``. + It drives the detached-launch / poll-marker handshake of + ``harness.common.run_agent`` (via ``sandbox.exec_and_wait``) without any real + process: when it sees the ``setsid`` launch command it awaits the injected + ``on_launch(env)`` agent coroutine, then writes its exit code into the + done-marker file so the next poll succeeds. + + Construct directly, or via :meth:`factory` to get a zero-arg callable that + ``examples...generate.E2BSandbox`` / ``swe.E2BSandbox`` can be monkeypatched + to (they call ``E2BSandbox(image)``). + """ + + def __init__( + self, + image: str = "fake-image", + *, + on_launch: Callable[[dict], Awaitable[int]] | None = None, + responses: list[tuple[str, tuple[int, str, str]]] | None = None, + ) -> None: + self.image = image + self.sandbox_id = f"fake-{image}" + self.on_launch = on_launch + # ordered (substring -> (exit, stdout, stderr)) overrides, first match wins. + self.responses = list(responses or []) + self.files: dict[str, str | bytes] = {} + self.exec_log: list[tuple[str, str]] = [] # (cmd, user) + + @classmethod + def factory(cls, **kwargs) -> Callable[..., FakeSandbox]: + """Return ``E2BSandbox(image)``-compatible constructor with kwargs baked in.""" + + def _make(image: str = "fake-image", **_ignored) -> FakeSandbox: + return cls(image, **kwargs) + + return _make + + async def __aenter__(self) -> FakeSandbox: + return self + + async def __aexit__(self, *exc) -> None: + return None + + async def exec(self, cmd, *, user="root", env=None, timeout=120, check=False, idempotent=True): + self.exec_log.append((cmd, user)) + + # Detached launch (run_agent): drive the fake agent, then drop the marker. + if "setsid" in cmd and self.on_launch is not None: + code = await self.on_launch(env or {}) + done = _done_path_from_launch(cmd) + if done: + self.files[done] = f"{code}\n" + return 0, "", "" + + # Marker poll (run_agent): succeed only once the marker file exists. + m = _POLL_RE.search(cmd) + if m: + path = m.group(1) + if path in self.files: + return 0, _as_str(self.files[path]), "" + return 1, "", "" + + for needle, result in self.responses: + if needle in cmd: + return result + return 0, "", "" + + async def write_file(self, sandbox_path, content, *, user="root") -> None: + self.files[sandbox_path] = content + + async def read_file(self, sandbox_path, *, user="root") -> str: + return _as_str(self.files.get(sandbox_path, "")) + + +def _as_str(v: str | bytes) -> str: + return v.decode() if isinstance(v, bytes) else v + + +def _done_path_from_launch(cmd: str) -> str | None: + """Recover the exit-code marker path from a ``setsid bash {launcher}`` command + so the subsequent poll matches. ``sandbox.exec_and_wait`` names the launcher + ``/tmp/.{tag}.sh`` and its sibling marker ``/tmp/.{tag}.done``.""" + m = re.search(r"setsid bash (\S+)\.sh\b", cmd) + if m: + return f"{m.group(1)}.done" + return None diff --git a/tests/test_agent/test_adapters.py b/tests/test_agent/test_adapters.py new file mode 100644 index 000000000..e7c6ec601 --- /dev/null +++ b/tests/test_agent/test_adapters.py @@ -0,0 +1,467 @@ +"""Unit tests for the agent HTTP adapters (Anthropic + OpenAI) and parsing. + +Every test drives a REAL adapter over a real aiohttp loopback +(``TestServer``/``TestClient``) and a real ``/inference/v1/generate`` upstream +(:class:`tests.test_agent._fakes.FakeVLLMServer`) -- so the whole +translate -> vllm -> parse -> record_turn -> finish_session path runs +unmocked; only the model server and tokenizer are faked. Covers both wire +protocols plus the standalone parsing helpers in ``vime.agent.parsing``. + +Replaces the pre-refactor ``tests/test_agent_adapters.py`` (which imported now- +removed symbols and a dropped ``/v1/responses`` endpoint). +""" + +from __future__ import annotations + +import asyncio +import json +import sys +from pathlib import Path + +import pytest +from aiohttp.test_utils import TestClient, TestServer + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from tests.test_agent._fakes import FakeTokenizer, FakeVLLMServer # noqa: E402 + +from vime.agent.adapters import anthropic, openai # noqa: E402 +from vime.agent.parsing import parse_model_output, parse_xml_tool_uses # noqa: E402 +from vime.utils.types import Sample # noqa: E402 + +NUM_GPUS = 0 + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +class _Headers: + def __init__(self, headers: dict[str, str]) -> None: + self.headers = headers + + +def _parse_sse(raw: str) -> list[tuple[str, object]]: + """Parse an SSE byte-stream body into ``(event_name, payload)`` pairs; + ``payload`` is the decoded JSON, or the literal ``"[DONE]"``.""" + events: list[tuple[str, object]] = [] + event_name = "message" + data_lines: list[str] = [] + + def flush() -> None: + nonlocal event_name, data_lines + if data_lines: + data = "\n".join(data_lines) + events.append((event_name, data if data == "[DONE]" else json.loads(data))) + event_name = "message" + data_lines = [] + + for line in raw.splitlines(): + if not line: + flush() + elif line.startswith("event:"): + event_name = line.removeprefix("event:").strip() + elif line.startswith("data:"): + data_lines.append(line.removeprefix("data:").strip()) + flush() + return events + + +async def _drain(adapter, sid) -> list[Sample]: + return await adapter.finish_session(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + + +# =========================================================================== +# §1 session-id resolution +# =========================================================================== + + +def test_anthropic_session_id_prefers_bearer_then_api_key(): + assert anthropic._request_session_id(_Headers({"X-Api-Key": "key"})) == "key" + assert anthropic._request_session_id(_Headers({"Authorization": "Bearer bsid", "X-Api-Key": "key"})) == "bsid" + assert anthropic._request_session_id(_Headers({})) == "default" + + +def test_openai_session_id_prefers_bearer_then_body(): + req = _Headers({}) + assert openai._request_session_id(req, {"metadata": {"session_id": "meta"}, "user": "u"}) == "meta" + assert openai._request_session_id(req, {"user": "u"}) == "u" + assert openai._request_session_id(_Headers({"Authorization": "Bearer bsid"}), {"user": "u"}) == "bsid" + assert openai._request_session_id(req, {}) == "default" + + +# =========================================================================== +# §2 translation (wire -> chat-template messages) +# =========================================================================== + + +def test_anthropic_translation_keeps_tool_results_thinking_and_tools(): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan"}, + {"type": "text", "text": "ok"}, + {"type": "tool_use", "name": "lookup", "input": {"q": "vime"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "u1", "content": "result"}]}, + ] + translated = anthropic._translate_messages(messages, system="sys") + assert translated == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "ok", + "reasoning_content": "plan", + "tool_calls": [{"type": "function", "function": {"name": "lookup", "arguments": {"q": "vime"}}}], + }, + {"role": "tool", "content": "result"}, + ] + tools = anthropic._tools_to_chat_tools( + [{"name": "lookup", "description": "search", "input_schema": {"type": "object"}}] + ) + assert tools == [ + {"type": "function", "function": {"name": "lookup", "description": "search", "parameters": {"type": "object"}}} + ] + + +def test_openai_translation_developer_to_system_and_tool_calls_to_dict(): + translated = openai._translate_messages( + [ + {"role": "developer", "content": "rules"}, + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": '{"q": "vime"}'}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + ) + assert translated == [ + {"role": "system", "content": "rules"}, + {"role": "user", "content": "hello"}, + # wire-only id dropped; arguments coerced JSON-string -> dict. + { + "role": "assistant", + "content": "", + "tool_calls": [{"type": "function", "function": {"name": "lookup", "arguments": {"q": "vime"}}}], + }, + # tool_call_id dropped. + {"role": "tool", "content": "found"}, + ] + + +# =========================================================================== +# §3 non-stream JSON + token capture (real HTTP, real /inference/v1/generate) +# =========================================================================== + + +def test_anthropic_messages_nonstream_records_token_segments(): + async def run_case(): + async with FakeVLLMServer([[(-0.1, 101), (-0.2, 102)]]) as vllm: + tok = FakeTokenizer(outputs={(101, 102): "done now"}) + adapter = anthropic.AnthropicAdapter(tokenizer=tok, vllm_url=vllm.url) + adapter.open_session("sid-a") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-a"}, + json={"model": "m", "max_tokens": 7, "messages": [{"role": "user", "content": "hi"}]}, + ) + data = await resp.json() + finally: + await client.close() + samples = await _drain(adapter, "sid-a") + + assert resp.status == 200 + assert data["type"] == "message" and data["stop_reason"] == "end_turn" + assert data["content"] == [{"type": "text", "text": "done now"}] + # adapter posted the rendered prompt ids and capped max_tokens at the request cap. + assert vllm.requests[0]["sampling_params"]["max_tokens"] == 7 + assert vllm.routing_keys == ["sid-a"] + # one trained turn: the two response ids carry loss=1 + real logprobs. + assert len(samples) == 1 + s = samples[0] + assert s.tokens[-2:] == [101, 102] + assert s.loss_mask[-2:] == [1, 1] + assert s.rollout_log_probs[-2:] == [-0.1, -0.2] + assert s.response == "done now" + + asyncio.run(run_case()) + + +def test_openai_chat_completions_nonstream_records_token_segments(): + async def run_case(): + async with FakeVLLMServer([[(-0.3, 201)]]) as vllm: + tok = FakeTokenizer(outputs={(201,): "hello"}) + adapter = openai.OpenAIAdapter(tokenizer=tok, vllm_url=vllm.url) + adapter.open_session("sid-o") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-o"}, + json={"model": "m", "max_tokens": 4, "messages": [{"role": "user", "content": "hi?"}]}, + ) + data = await resp.json() + finally: + await client.close() + samples = await _drain(adapter, "sid-o") + + assert resp.status == 200 + assert data["object"] == "chat.completion" + assert data["choices"][0]["message"] == {"role": "assistant", "content": "hello"} + assert data["choices"][0]["finish_reason"] == "stop" + assert vllm.requests[0]["sampling_params"]["max_tokens"] == 4 + assert len(samples) == 1 and samples[0].tokens[-1] == 201 and samples[0].loss_mask[-1] == 1 + + asyncio.run(run_case()) + + +# =========================================================================== +# §4 streaming SSE +# =========================================================================== + + +def test_anthropic_messages_streams_blocks(): + async def run_case(): + async with FakeVLLMServer([[(-0.1, 301)]]) as vllm: + tok = FakeTokenizer(outputs={(301,): "streamed"}) + adapter = anthropic.AnthropicAdapter(tokenizer=tok, vllm_url=vllm.url) + adapter.open_session("sid-as") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-as", "Accept": "text/event-stream"}, + json={ + "model": "m", + "stream": True, + "max_tokens": 8, + "messages": [{"role": "user", "content": "x"}], + }, + ) + raw = await resp.text() + finally: + await client.close() + await _drain(adapter, "sid-as") + + names = [name for name, _ in _parse_sse(raw)] + assert names[0] == "message_start" + assert "content_block_delta" in names + assert names[-1] == "message_stop" + deltas = [p for n, p in _parse_sse(raw) if n == "content_block_delta"] + assert any(d["delta"].get("text") == "streamed" for d in deltas) + + asyncio.run(run_case()) + + +def test_openai_chat_completions_streams_chunks_until_done(): + async def run_case(): + async with FakeVLLMServer([[(-0.1, 401)]]) as vllm: + tok = FakeTokenizer(outputs={(401,): "streamed text"}) + adapter = openai.OpenAIAdapter(tokenizer=tok, vllm_url=vllm.url) + adapter.open_session("sid-os") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-os"}, + json={"model": "m", "stream": True, "messages": [{"role": "user", "content": "x"}]}, + ) + raw = await resp.text() + finally: + await client.close() + await _drain(adapter, "sid-os") + + events = _parse_sse(raw) + chunks = [p for _, p in events if isinstance(p, dict)] + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant"} + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert events[-1] == ("message", "[DONE]") + + asyncio.run(run_case()) + + +# =========================================================================== +# §5 multi-turn token alignment (tool call -> tool result -> answer) +# =========================================================================== + + +def test_anthropic_multiturn_wire_roundtrip_and_token_capture(): + """Two-turn round-trip: a tool-call turn, then a tool-result + answer turn. + + Asserts the adapter-level behaviour the branching test can't see: the wire + tool_use block round-trips, both turns route to the same sid, and + finish_session yields aligned training samples. (The fine-grained + clean/drift linearization is owned by test_trajectory_manager_branching.)""" + + async def run_case(): + r1 = "vime" + async with FakeVLLMServer([[(-0.5, 700), (-0.5, 701)], [(-0.4, 800)]]) as vllm: + tok = FakeTokenizer(outputs={(700, 701): r1, (800,): "the answer"}) + adapter = anthropic.AnthropicAdapter(tokenizer=tok, vllm_url=vllm.url) + adapter.open_session("sid-mt") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + tools = [ + {"name": "lookup", "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}} + ] + try: + first = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-mt"}, + json={ + "model": "m", + "max_tokens": 5, + "tools": tools, + "messages": [{"role": "user", "content": [{"type": "text", "text": "find vime"}]}], + }, + ) + fdata = await first.json() + tool_use = next(b for b in fdata["content"] if b["type"] == "tool_use") + second = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-mt"}, + json={ + "model": "m", + "max_tokens": 7, + "tools": tools, + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "find vime"}]}, + {"role": "assistant", "content": fdata["content"]}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": tool_use["id"], "content": "found"} + ], + }, + ], + }, + ) + await second.json() + finally: + await client.close() + samples = await _drain(adapter, "sid-mt") + + assert first.status == 200 and second.status == 200 + assert fdata["stop_reason"] == "tool_use" + assert tool_use["name"] == "lookup" and tool_use["input"] == {"query": "vime"} + # both turns routed to the same sid; the adapter posted the growing prompt. + assert vllm.routing_keys == ["sid-mt", "sid-mt"] + assert vllm.requests[1]["token_ids"][: len(vllm.requests[0]["token_ids"])] == vllm.requests[0]["token_ids"] + # finish_session produces at least one aligned, partly-trained sample. + assert samples + for s in samples: + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length + assert sum(s.loss_mask) > 0 + + asyncio.run(run_case()) + + +# =========================================================================== +# §6 adapter behaviour: turn cap, mid-list system fold +# =========================================================================== + + +def test_max_turns_per_sid_returns_429(): + async def run_case(): + async with FakeVLLMServer([[(-0.1, 501)], [(-0.1, 502)]]) as vllm: + tok = FakeTokenizer() + adapter = anthropic.AnthropicAdapter(tokenizer=tok, vllm_url=vllm.url, max_turns_per_sid=1) + adapter.open_session("sid-cap") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + body = {"model": "m", "max_tokens": 4, "messages": [{"role": "user", "content": "x"}]} + h = {"Authorization": "Bearer sid-cap"} + first = await client.post("/v1/messages", headers=h, json=body) + second = await client.post("/v1/messages", headers=h, json=body) + finally: + await client.close() + await _drain(adapter, "sid-cap") + assert first.status == 200 + assert second.status == 429 + + asyncio.run(run_case()) + + +def test_mid_list_system_folds_into_user(): + body = { + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "system", "content": "skills list"}, + {"role": "user", "content": "next"}, + ] + } + changed = anthropic._fold_mid_list_system_into_user(body) + assert changed + # the mid-list system message is gone; its text is wrapped into the prior user. + assert [m["role"] for m in body["messages"]] == ["user", "user"] + folded = body["messages"][0]["content"] + assert any(b.get("text", "").startswith("") for b in folded) + + +# =========================================================================== +# §7 parsing helpers (vime.agent.parsing) +# =========================================================================== + + +def test_parse_model_output_plain_text_no_parsers(): + # tokenizer is only used on the parser paths (#198 made it required for vLLM + # parsers); the no-parser passthrough never touches it, so None is fine here. + parsed = parse_model_output( + "just text", tokenizer=None, tools_schema=None, tool_parser_name=None, reasoning_parser_name=None + ) + assert parsed.text == "just text" + assert parsed.tool_uses == [] + assert parsed.reasoning == "" + + +def test_parse_model_output_think_split_fallback(): + # The qwen3 reasoning parser lives in vllm (lazy import); skip where the + # lean CPU CI env has no vllm installed. + pytest.importorskip("vllm.entrypoints.openai.chat_completion.protocol") + parsed = parse_model_output( + "reason herevisible", + tools_schema=None, + tool_parser_name=None, + reasoning_parser_name="qwen3", + ) + # the qwen3 reasoning parser (or the fallback) splits reasoning out. + assert "visible" in parsed.text + assert "reason here" in parsed.reasoning + + +def test_parse_xml_tool_uses_fallback(): + raw = "lead vime tail" + schema = [{"function": {"name": "lookup"}}] + cleaned, uses = parse_xml_tool_uses(raw, schema) + assert uses == [{"name": "lookup", "input": {"q": "vime"}}] + assert "" not in cleaned + assert "lead" in cleaned and "tail" in cleaned + + +def test_parse_xml_tool_uses_ignores_unknown_tool(): + raw = "x" + cleaned, uses = parse_xml_tool_uses(raw, [{"function": {"name": "lookup"}}]) + assert uses == [] + assert "" in cleaned # left untouched + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_agent/test_agent_rollout_cpu.py b/tests/test_agent/test_agent_rollout_cpu.py new file mode 100644 index 000000000..66bda4c99 --- /dev/null +++ b/tests/test_agent/test_agent_rollout_cpu.py @@ -0,0 +1,292 @@ +"""CPU-only agent rollout test: the whole pipeline, no GPU / E2B / vllm. + +This walks a real agent rollout end to end on CPU. Only four external edges are +faked (see ``tests/test_agent/_fakes.py``): the tokenizer, the E2B sandbox, the +vllm ``/inference/v1/generate`` server, and the agent CLI process inside the sandbox. +Everything between -- ``generate.generate`` orchestration, the in-thread adapter +HTTP app, wire translation, ``record_turn`` / tree building, ``finish_session`` +linearization, ``swe`` workspace-prep / diff / eval, the harness lifecycle and +its detached-launch transport -- is the real code. + +The "agent" is a coroutine standing in for ``claude -p`` / ``codex exec``: the +sandbox fake invokes it on launch, and it dials the adapter back over real HTTP +loopback (``trust_env=False`` so the cluster proxy can't hijack 127.0.0.1), +firing a couple of turns the way the real CLI would. + +Two protocol chains are covered: + + * ``test_generate_*`` -- the production path: real ``generate.generate()``, + which is hardwired to ClaudeCodeHarness + AnthropicAdapter. + * ``test_codex_openai_rollout_closes_loop`` -- the same loop for the + CodexHarness + OpenAIAdapter pair, hand-wired (generate() does not select it). +""" + +from __future__ import annotations + +import asyncio +import contextlib +import dataclasses +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import aiohttp +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +# Importing generate pulls vime.utils.processing_utils at module load, which +# eagerly imports transformers -- a heavy dep deliberately absent from the +# CPU-only CI env for this test. We never touch a real tokenizer (load_tokenizer +# is patched with FakeTokenizer below), so stub transformers before the import +# so the chain resolves without it. +if "transformers" not in sys.modules: + _tf_stub = types.ModuleType("transformers") + for _name in ("AutoProcessor", "AutoTokenizer", "PreTrainedTokenizerBase", "ProcessorMixin"): + setattr(_tf_stub, _name, type(_name, (), {})) + sys.modules["transformers"] = _tf_stub + +# generate.generate() uses asyncio.timeout(), a 3.11+ API. CI runs the agent +# tests on 3.10, so shim it onto wait_for. The wall-clock guard never fires in +# these tests (every case finishes well under the guard), so a thin pass-through +# context manager is enough. +if not hasattr(asyncio, "timeout"): + + @contextlib.asynccontextmanager + async def _timeout_shim(_delay): + yield + + asyncio.timeout = _timeout_shim + +import examples.coding_agent_rl.generate as gen # noqa: E402 +import examples.coding_agent_rl.swe as swe # noqa: E402 +from tests.test_agent._fakes import FakeSandbox, FakeTokenizer, fake_call_vllm_generate # noqa: E402 + +from vime.agent.adapters import OpenAIAdapter # noqa: E402 +from vime.agent.adapters import common as adapters_common # noqa: E402 +from vime.agent.aiohttp_threaded import run_app_in_thread # noqa: E402 +from vime.agent.harness import ClaudeCodeHarness, CodexHarness # noqa: E402 +from vime.agent.harness import common as harness_common # noqa: E402 +from vime.utils.misc import SingletonMeta # noqa: E402 +from vime.utils.types import Sample # noqa: E402 + +NUM_GPUS = 0 + +_REAL_SLEEP = asyncio.sleep + + +async def _fast_sleep(_secs): # collapse run_command's 5s poll loop + await _REAL_SLEEP(0) + + +def _args() -> SimpleNamespace: + return SimpleNamespace( + hf_checkpoint="unused", # load_tokenizer is patched + rollout_max_context_len=0, + vllm_tool_call_parser=None, + vllm_reasoning_parser=None, + vllm_router_ip="127.0.0.1", + vllm_router_port=1, # never dialed (call_vllm_generate is patched) + ) + + +def _base_sample(**md) -> Sample: + meta = { + "instance_id": "demo-1", + "image": "fake-image", + "workdir": "/workspace/repo", + "problem_statement": "fix the bug", + "eval_cmd": "true", + **md, + } + return Sample(index=0, group_index=0, prompt="fix the bug", metadata=meta) + + +async def _anthropic_agent(env: dict, *, n_turns: int = 2) -> int: + """Stand-in for ``claude -p``: dial the adapter back over real HTTP and fire + a few Anthropic turns, reading wiring from the env the harness exported.""" + base_url = env["ANTHROPIC_BASE_URL"] + token = env["ANTHROPIC_AUTH_TOKEN"] + history = [{"role": "user", "content": [{"type": "text", "text": "solve the issue"}]}] + async with aiohttp.ClientSession(trust_env=False) as sess: + for _ in range(n_turns): + async with sess.post( + f"{base_url}/v1/messages", + headers={"Authorization": f"Bearer {token}"}, + json={"model": "m", "max_tokens": 64, "messages": history}, + ) as r: + data = await r.json() + history.append({"role": "assistant", "content": data["content"]}) + history.append({"role": "user", "content": [{"type": "text", "text": "continue"}]}) + return 0 + + +def _patch_generate(monkeypatch, tokenizer: FakeTokenizer, sandbox_factory) -> None: + """Wire generate.generate()'s four external edges to CPU fakes.""" + # in-thread adapter must bind to loopback and need no public host check. + monkeypatch.setattr( + gen, + "CONFIG", + dataclasses.replace( + gen.CONFIG, + adapter_public_host="127.0.0.1", + adapter_bind_host="127.0.0.1", + adapter_port=0, + rollout_guard_sec=60, + agent_time_budget_sec=30, + eval_timeout_sec=30, + boot_retries=1, + ), + ) + monkeypatch.setattr(gen, "load_tokenizer", lambda *a, **k: tokenizer) + monkeypatch.setattr(gen, "E2BSandbox", sandbox_factory) # boot sandbox + monkeypatch.setattr(swe, "E2BSandbox", sandbox_factory) # eval sandbox + monkeypatch.setattr(ClaudeCodeHarness, "install_cli", _noop_install) + monkeypatch.setattr(harness_common.asyncio, "sleep", _fast_sleep) + monkeypatch.setattr(adapters_common, "call_vllm_generate", fake_call_vllm_generate(_two_turn_script(), tokenizer)) + # _AdapterService is a SingletonMeta singleton; drop any cached instance so + # each test builds a fresh adapter + app thread. + SingletonMeta.clear_instances(gen._AdapterService) + + +async def _noop_install(self, sb) -> None: + return None + + +def _two_turn_script(): + # (response_text, finish_reason, logprobs) per vllm call; encoded by the + # FakeTokenizer so the adapter's decode round-trips it. + return [ + ("let me look at the code", "stop", None), + ("the fix is applied done", "stop", None), + ] + + +# =========================================================================== +# §1 production path: real generate() over ClaudeCode + Anthropic +# =========================================================================== + + +def test_generate_produces_trained_samples(): + async def run_case(monkeypatch): + tok = FakeTokenizer() + sandbox_factory = FakeSandbox.factory(on_launch=_anthropic_agent) + _patch_generate(monkeypatch, tok, sandbox_factory) + + samples = await gen.generate(_args(), _base_sample(), sampling_params={"max_new_tokens": 32}) + + assert samples, "rollout produced no samples" + for s in samples: + assert s.status == Sample.Status.COMPLETED + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length + assert sum(s.loss_mask) > 0 # at least one trained token + assert s.metadata.get("agent_exit_code") == 0 + # eval_cmd "true" applied cleanly on a clean (empty) diff -> reward 1.0, + # split evenly across the emitted samples. + assert abs(sum(s.reward for s in samples) - 1.0) < 1e-9 + + with pytest.MonkeyPatch.context() as mp: + asyncio.run(run_case(mp)) + + +def test_generate_aborts_on_empty_trajectory(): + """If the agent never drives a turn, the session is empty and generate() + returns a single ABORTED sample (the fan-out shape) rather than crashing.""" + + async def silent_agent(_env) -> int: + return 0 # never contacts the adapter -> empty trajectory + + async def run_case(monkeypatch): + tok = FakeTokenizer() + sandbox_factory = FakeSandbox.factory(on_launch=silent_agent) + _patch_generate(monkeypatch, tok, sandbox_factory) + + samples = await gen.generate(_args(), _base_sample(), sampling_params={}) + + assert len(samples) == 1 + assert samples[0].status == Sample.Status.ABORTED + assert samples[0].metadata.get("abort_reason") == "adapter_session_empty" + + with pytest.MonkeyPatch.context() as mp: + asyncio.run(run_case(mp)) + + +def test_generate_aborts_on_missing_image(): + async def run_case(monkeypatch): + tok = FakeTokenizer() + _patch_generate(monkeypatch, tok, FakeSandbox.factory(on_launch=_anthropic_agent)) + # blank image -> early abort before any sandbox boot. + samples = await gen.generate(_args(), _base_sample(image=""), sampling_params={}) + assert len(samples) == 1 + assert samples[0].status == Sample.Status.ABORTED + assert samples[0].metadata.get("abort_reason") == "missing_image_or_workdir" + + with pytest.MonkeyPatch.context() as mp: + asyncio.run(run_case(mp)) + + +# =========================================================================== +# §2 the Codex + OpenAI pair closes the same loop (hand-wired) +# =========================================================================== + + +async def _codex_agent(env: dict, *, n_turns: int = 2) -> int: + base_url = env["OPENAI_BASE_URL"] # already includes /v1 + token = env["OPENAI_API_KEY"] + history = [{"role": "user", "content": "solve the issue"}] + async with aiohttp.ClientSession(trust_env=False) as sess: + for _ in range(n_turns): + async with sess.post( + f"{base_url}/chat/completions", + headers={"Authorization": f"Bearer {token}"}, + json={"model": "m", "max_tokens": 64, "messages": history}, + ) as r: + data = await r.json() + msg = data["choices"][0]["message"] + history.append({"role": "assistant", "content": msg.get("content") or ""}) + history.append({"role": "user", "content": "continue"}) + return 0 + + +def test_codex_openai_rollout_closes_loop(monkeypatch): + """CodexHarness drives an in-thread OpenAIAdapter through a FakeSandbox; the + loop produces trained samples just like the production Anthropic path.""" + + async def run_case(): + tok = FakeTokenizer() + monkeypatch.setattr(adapters_common, "call_vllm_generate", fake_call_vllm_generate(_two_turn_script(), tok)) + monkeypatch.setattr(harness_common.asyncio, "sleep", _fast_sleep) + + adapter = OpenAIAdapter(tokenizer=tok, vllm_url="http://unused") + handle = run_app_in_thread(adapter.app, host="127.0.0.1", port=0, thread_name="test-openai-adapter") + adapter_url = f"http://127.0.0.1:{handle.port}" + sid = "codex-sess" + adapter.open_session(sid) + try: + sb = FakeSandbox(on_launch=_codex_agent) + rc = await CodexHarness().run( + sb, + workdir="/workspace/repo", + session_id=sid, + adapter_url=adapter_url, + time_budget_sec=30, + prompt="fix it", + ) + samples = await adapter.finish_session(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + finally: + handle.stop() + + assert rc == 0 + assert samples + for s in samples: + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length + assert sum(s.loss_mask) > 0 + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_agent/test_harness.py b/tests/test_agent/test_harness.py new file mode 100644 index 000000000..6335e9abc --- /dev/null +++ b/tests/test_agent/test_harness.py @@ -0,0 +1,236 @@ +"""Unit tests for the coding-agent harness + sandbox layers. + +These cover the parts a happy-path rollout can't pin down precisely: that each +harness writes the right CLI config and launches with the right command + env, +that ``run_agent``'s detached-launch / poll-marker handshake returns the right +exit code (and times out correctly), and that ``ensure_agent_user`` issues the +expected provisioning command. A :class:`tests.test_agent._fakes.FakeSandbox` +records every ``exec`` / ``write_file`` so we assert on the issued commands +without a real sandbox or any root privilege. +""" + +from __future__ import annotations + +import asyncio +import base64 +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from tests.test_agent._fakes import FakeSandbox # noqa: E402 + +from vime.agent import sandbox as sandbox_mod # noqa: E402 +from vime.agent.harness import ClaudeCodeHarness, CodexHarness, HarnessContext # noqa: E402 +from vime.agent.harness import common as hc # noqa: E402 + +NUM_GPUS = 0 + +# Run the 5s poll loop instantly without recursing into the patched function. +_REAL_SLEEP = asyncio.sleep + + +async def _fast_sleep(_secs): + await _REAL_SLEEP(0) + + +def _ctx(workdir="/workspace/repo", sid="sess-1", url="http://host:18001") -> HarnessContext: + return HarnessContext(workdir=workdir, session_id=sid, adapter_url=url) + + +def _find(exec_log, needle): + return [cmd for cmd, _user in exec_log if needle in cmd] + + +# =========================================================================== +# §1 run_agent handshake (the E2B detached-launch transport) +# =========================================================================== + + +def test_run_agent_returns_marker_exit_code(): + async def run_case(): + seen = {} + + async def fake_agent(env): + seen["env"] = env + return 0 # exit code written into the done marker + + sb = FakeSandbox(on_launch=fake_agent) + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await hc.run_agent( + sb, workdir="/workspace/repo", start_cmd="claude -p hi", env={"A": "1"}, time_budget_sec=30 + ) + assert rc == 0 + assert seen["env"] == {"A": "1"} + # launcher script + detached setsid launch all issued, exit code captured. + assert any("run.sh" in p for p in sb.files) + assert _find(sb.exec_log, "setsid") + assert any("echo $?" in v for v in sb.files.values()) + + asyncio.run(run_case()) + + +def test_run_agent_propagates_nonzero_exit(): + async def run_case(): + async def fail_agent(_env): + return 7 + + sb = FakeSandbox(on_launch=fail_agent) + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await hc.run_agent(sb, workdir="/w", start_cmd="x", env={}, time_budget_sec=30) + assert rc == 7 + + asyncio.run(run_case()) + + +def test_run_agent_times_out_when_marker_never_appears(): + async def run_case(): + sb = FakeSandbox(on_launch=None) # no agent -> marker never written + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await hc.run_agent(sb, workdir="/w", start_cmd="x", env={}, time_budget_sec=0) + assert rc == sandbox_mod.EXIT_TIME_BUDGET_EXCEEDED + + asyncio.run(run_case()) + + +# =========================================================================== +# §2 ClaudeCodeHarness config + launch +# =========================================================================== + + +def test_claude_code_write_config_preacks_bypass_permissions(): + async def run_case(): + sb = FakeSandbox() + await ClaudeCodeHarness().write_config(sb, _ctx()) + joined = " ".join(cmd for cmd, _ in sb.exec_log) + assert "/home/agent/.claude/settings.json" in joined + assert "bypassPermissionsModeAccepted" in joined + assert "hasCompletedOnboarding" in joined + + asyncio.run(run_case()) + + +def test_claude_code_launch_command_and_env(): + async def run_case(): + captured = {} + + async def agent(env): + captured["env"] = env + return 0 + + capturing = FakeSandbox(on_launch=agent) + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await ClaudeCodeHarness().launch_and_wait( + capturing, _ctx(sid="sess-cc", url="http://host:18001"), prompt="solve it", time_budget_sec=30 + ) + assert rc == 0 + # the prompt + flags land in the launcher script body. + body = next(v for k, v in capturing.files.items() if k.endswith("run.sh")) + assert "claude -p 'solve it'" in body + assert "--permission-mode bypassPermissions" in body + # env carries the adapter wiring under the Anthropic var names. + env = captured["env"] + assert env["ANTHROPIC_BASE_URL"] == "http://host:18001" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sess-cc" + assert env["ANTHROPIC_MODEL"] == "vime-actor" + assert env["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + asyncio.run(run_case()) + + +# =========================================================================== +# §3 CodexHarness config + launch +# =========================================================================== + + +def test_codex_write_config_base64_roundtrips_inline_base_url(): + async def run_case(): + sb = FakeSandbox() + await CodexHarness().write_config(sb, _ctx(url="http://host:18001")) + # config written via base64 round-trip; decode the captured payload. + cmd = next(c for c, _ in sb.exec_log if "base64 -d > /home/agent/.codex/config.toml" in c) + b64 = cmd.split("echo ")[1].split(" | base64")[0].strip("'") + toml = base64.b64decode(b64).decode() + assert 'base_url = "http://host:18001/v1"' in toml # MUST be inline + assert 'wire_api = "chat"' in toml + assert 'model_provider = "vime"' in toml + + asyncio.run(run_case()) + + +def test_codex_launch_command_and_env(): + async def run_case(): + captured = {} + + async def agent(env): + captured["env"] = env + return 0 + + sb = FakeSandbox(on_launch=agent) + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await CodexHarness().launch_and_wait( + sb, _ctx(sid="sess-cx", url="http://host:18001"), prompt="do work", time_budget_sec=30 + ) + assert rc == 0 + body = next(v for k, v in sb.files.items() if k.endswith("run.sh")) + assert "codex exec" in body and "do work" in body and "--skip-git-repo-check" in body + env = captured["env"] + assert env["OPENAI_API_KEY"] == "sess-cx" + assert env["OPENAI_BASE_URL"] == "http://host:18001/v1" + + asyncio.run(run_case()) + + +# =========================================================================== +# §4 ensure_agent_user (sandbox infra) +# =========================================================================== + + +def test_ensure_agent_user_provisions_user_and_git_safe_dir(): + async def run_case(): + sb = FakeSandbox() + await sandbox_mod.ensure_agent_user(sb, "/workspace/repo") + cmd = next(c for c, _ in sb.exec_log if "useradd" in c) + assert "id agent" in cmd + assert "chown -R agent:agent" in cmd and "/workspace/repo" in cmd + assert "git config --system --add safe.directory '*'" in cmd + + asyncio.run(run_case()) + + +# =========================================================================== +# §5 harness.run wires the steps in order +# =========================================================================== + + +def test_base_harness_run_calls_steps_in_order(): + async def run_case(): + async def agent(_env): + return 0 + + sb = FakeSandbox(on_launch=agent) + with patch.object(hc.asyncio, "sleep", new=_fast_sleep): + rc = await ClaudeCodeHarness().run( + sb, + workdir="/workspace/repo", + session_id="sess-run", + adapter_url="http://host:18001", + time_budget_sec=30, + prompt="go", + ) + assert rc == 0 + joined = " ".join(c for c, _ in sb.exec_log) + # ensure_agent_user (useradd) -> write_config (settings.json) -> launch (setsid) + order = [k for k in ("useradd", "settings.json", "setsid") if k in joined] + assert order == ["useradd", "settings.json", "setsid"] + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_agent/test_sandbox_exec_and_wait.py b/tests/test_agent/test_sandbox_exec_and_wait.py new file mode 100644 index 000000000..a5cfc0446 --- /dev/null +++ b/tests/test_agent/test_sandbox_exec_and_wait.py @@ -0,0 +1,167 @@ +"""CPU tests for ``vime.agent.sandbox.exec_and_wait``'s spawn-lock contract. + +The detached spawn is guarded by ``mkdir {lock_dir} || exit 0`` so that a +transport-level retry of the *same* spawn RPC (a severed response replayed by +``E2BSandbox._rpc_retry``) cannot double-execute the command. That guard must +not leak into the next *logical* invocation of the same tag: the lock dir used +to survive forever, so a second ``exec_and_wait`` with the same tag skipped the +spawn at the guard (before the stale-marker cleanup, which sat behind it) and +``_await_done_marker`` immediately read the previous run's exit-code marker. + +Concrete victim: ``harness.common.install_npm_cli`` retries a failed npm +install three times with ``tag="harness-npm-install"`` — attempts 2 and 3 ran +nothing and returned attempt 1's exit code and log verbatim. + +The fake sandbox here interprets the actual command strings ``exec_and_wait`` +issues (mkdir guard short-circuit, rm cleanup, setsid launch, marker polls), so +these tests hold for any command layout that keeps the semantics right and fail +for one that reuses stale state. +""" + +from __future__ import annotations + +import asyncio +import re +import shlex +from types import SimpleNamespace + +import pytest + +import vime.agent.sandbox as sandbox_mod +from vime.agent.sandbox import exec_and_wait + +NUM_GPUS = 0 + + +_POLL_RE = re.compile(r"test -f (\S+) && cat \1") +_SPAWN_RE = re.compile(r"mkdir (\S+) 2>/dev/null \|\| exit 0; (.*)$") +_LAUNCH_RE = re.compile(r"setsid bash (\S+) ") +_TAIL_RE = re.compile(r"tail -c \d+ (\S+)") + + +class ShellFakeSandbox: + """Interprets exec_and_wait's shell commands against an in-memory FS. + + ``run_script`` is called once per *actual* launch with the launcher path; + it returns ``(exit_code, output)`` which land in the done/out marker files + exactly like the real detached command would write them. + """ + + sandbox_id = "shell-fake" + + def __init__(self, run_script): + self.run_script = run_script + self.files: dict[str, str] = {} + self.dirs: set[str] = set() + self.launches = 0 + self.exec_log: list[str] = [] + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return None + + async def write_file(self, path, content, *, user="root"): + self.files[path] = content + + async def read_file(self, path, *, user="root"): + return self.files.get(path, "") + + async def exec(self, cmd, *, user="root", env=None, timeout=120, check=False, idempotent=True): + self.exec_log.append(cmd) + + poll = _POLL_RE.search(cmd) + if poll: + path = poll.group(1) + if path in self.files: + return 0, self.files[path], "" + return 1, "", "" + + tail = _TAIL_RE.search(cmd) + if tail: + return 0, self.files.get(tail.group(1), "")[-512:], "" + + spawn = _SPAWN_RE.search(cmd) + if spawn: + lock_dir, rest = spawn.groups() + if lock_dir in self.dirs: + return 0, "", "" # guard hit: nothing after it runs + self.dirs.add(lock_dir) + self._run_shell_fragment(rest) + return 0, "", "" + + # Plain cleanup command(s): rm -rf / rm -f sequences. + self._run_shell_fragment(cmd) + return 0, "", "" + + def _run_shell_fragment(self, fragment): + for part in fragment.split(";"): + part = part.strip().rstrip("&").strip() + if part.startswith(("rm -rf", "rm -f")): + for token in shlex.split(part)[2:]: + self.files.pop(token, None) + self.dirs.discard(token) + launch = _LAUNCH_RE.search(part + " ") + if launch: + launcher = launch.group(1) + assert launcher in self.files, "launcher must be written before the spawn" + self.launches += 1 + exit_code, output = self.run_script(self.launches) + out_file = re.search(r"> (\S+) 2>&1", part).group(1) + done_file = launcher.replace(".sh", ".done") + self.files[out_file] = output + self.files[done_file] = f"{exit_code}\n" + + +@pytest.fixture +def fast_marker_polls(monkeypatch): + """_await_done_marker sleeps 5s between polls; make that instant.""" + + async def _instant(_seconds): + return None + + monkeypatch.setattr(sandbox_mod, "asyncio", SimpleNamespace(sleep=_instant)) + + +@pytest.mark.unit +def test_same_tag_reinvocation_actually_reruns(fast_marker_polls): + """A retry loop (e.g. install_npm_cli) must re-run the command, not be fed + the previous attempt's stale exit code and log.""" + outcomes = {1: (1, "attempt-1 failed"), 2: (0, "attempt-2 ok")} + sb = ShellFakeSandbox(run_script=lambda n: outcomes[n]) + + async def _two_attempts(): + first = await exec_and_wait(sb, cmd="npm install", time_budget_sec=60, tag="npm", want_output=True) + second = await exec_and_wait(sb, cmd="npm install", time_budget_sec=60, tag="npm", want_output=True) + return first, second + + (first_code, first_out), (second_code, second_out) = asyncio.run(_two_attempts()) + + assert (first_code, first_out) == (1, "attempt-1 failed") + assert sb.launches == 2, "second invocation must actually spawn the command" + assert (second_code, second_out) == (0, "attempt-2 ok") + + +@pytest.mark.unit +def test_transport_retry_of_the_spawn_stays_deduped(fast_marker_polls): + """Replaying the spawn RPC itself (what the mkdir guard is *for*) must not + double-execute the command.""" + sb = ShellFakeSandbox(run_script=lambda n: (0, "ok")) + + asyncio.run(exec_and_wait(sb, cmd="true", time_budget_sec=60, tag="job")) + + spawn_cmds = [c for c in sb.exec_log if "setsid" in c] + assert len(spawn_cmds) == 1 + # Replay the identical spawn RPC, as _rpc_retry would after a severed + # response: the guard must swallow it. + asyncio.run(sb.exec(spawn_cmds[0])) + assert sb.launches == 1 + + # The per-invocation cleanup must NOT ride inside the guarded spawn — + # behind the guard it never runs on a replayed tag. + assert not any("rm -" in c and "setsid" in c for c in sb.exec_log) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent/test_trajectory_manager_branching.py b/tests/test_agent/test_trajectory_manager_branching.py new file mode 100644 index 000000000..878f357e5 --- /dev/null +++ b/tests/test_agent/test_trajectory_manager_branching.py @@ -0,0 +1,1394 @@ +"""Branching-matrix tests for TrajectoryManager via record_turn / get_trajectory. + +This script drives the two public interfaces of +``vime.agent.trajectory.TrajectoryManager`` and exhaustively covers the +ways a trajectory can branch, organized as a two-axis matrix: + + * LAYER 1 — routing tree (record_turn). DFS merges on (role, message-equality) + only, so MESSAGE IDENTITY determines tree shape; token ids are irrelevant here. + * LAYER 2 — linearization (get_trajectory). TOKEN-ID prefix determines how each leaf + chain becomes Samples (clean continuation / drift case A·B1·B2 / cross-leaf + dedup / reward split). + * COMBINED — both layers interacting (rewrite-merge, tree-fork + token-drift + stacked, deep multi-leaf dedup, long mixed session). + +Readability: + Token ids are SEMANTIC small integers (see TOKEN_NAMES). Each message renders + to ``[START, ...body, END]`` with a per-role band, so an id like 2001 reads as + ``u:compute`` and 7001 reads as ````. Expected token sequences are built + with the same render_* helpers used to feed record_turn, never hand-typed + magic numbers. + +Dual mode: + Every case is a ``test_*`` function doing strict assertions, run under pytest + by default. Setting ``TRAJ_DUMP=1`` instead runs them via ``main()``, which + after each case prints the routing tree (token ids decoded to names) and every + linearized Sample with token / loss_mask aligned, so a human can read exactly + where each branch happened:: + + python tests/test_agent/test_trajectory_manager_branching.py # pytest + TRAJ_DUMP=1 python tests/test_agent/test_trajectory_manager_branching.py # human-readable dump +""" + +from __future__ import annotations + +import dataclasses # noqa: E402 +import os # noqa: E402 +import sys # noqa: E402 +from pathlib import Path # noqa: E402 + +import pytest # noqa: E402 + +# Run as a plain script (CI does `python `): make the repo root importable +# so `from tests.test_agent...` resolves without an installed `tests` package. +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from tests.test_agent._dump_helpers import dump_tree_txt # noqa: E402 + +from vime.agent.adapters.common import TurnRecord # noqa: E402 +from vime.agent.trajectory import TrajectoryManager, _common_prefix_len # noqa: E402 +from vime.utils.types import Sample # noqa: E402 + +# =========================================================================== +# §1 Semantic token vocabulary + reverse table +# =========================================================================== +# +# Per-role band. A message renders to [START, ...body, END]; the generation +# prompt appends the assistant START as the open-turn marker. + +_BANDS = { + "system": 1000, + "user": 2000, + "assistant": 9000, + "tool": 3000, +} +_GEN = _BANDS["assistant"] # add_generation_prompt marker +_DRIFT_BAND = 7000 + +# Reverse table: token id -> human-readable name. Filled lazily as messages are +# registered so dumps translate ids back to labels. +TOKEN_NAMES: dict[int, str] = {} +_ABBR = {"system": "sys", "user": "usr", "assistant": "ast", "tool": "tul"} +for _role, _base in _BANDS.items(): + TOKEN_NAMES[_base] = f"<{_ABBR[_role]}>" + TOKEN_NAMES[_base + 9] = f"" +TOKEN_NAMES[_GEN] = "" + + +def name_of(tok: int) -> str: + """Human-readable name for a token id (falls back to the raw int).""" + return TOKEN_NAMES.get(tok, str(tok)) + + +def _vis(label: str) -> str: + """Make whitespace visible in a token label for the dump. + + Whitespace-only drift (e.g. a trailing space from a cc rewrite) is invisible + in a terminal, which makes ``r:ok`` vs ``r:ok `` indistinguishable. Render + spaces as ``␣`` so the difference is obvious in the readable output. + """ + return label.replace(" ", "␣") + + +_ASST_BODY: dict[str, int] = {} + + +def _asst_body(label: str) -> int: + """Stable assistant body token for a response/message label. + + An assistant message replayed in a later prompt must render to the SAME + tokens the model generated for it, otherwise a clean continuation can never + hold (the cumulative prompt+response would not prefix the next prompt). So + both ``render_response`` and an assistant ``MsgTok`` derive their body token + from this one function, keyed on the label. Bodies are assigned by a stable + per-label counter (NOT a hash) so distinct labels never collide on one id — + a collision would mislabel tokens in the dump and could spuriously match + across turns. + """ + if label not in _ASST_BODY: + body = _BANDS["assistant"] + 100 + len(_ASST_BODY) + _ASST_BODY[label] = body + TOKEN_NAMES[body] = f"r:{_vis(label)}" + return _ASST_BODY[label] + + +def render_ids(ids: list[int]) -> str: + """Decode an id list into a space-joined readable string.""" + return " ".join(name_of(t) for t in ids) + + +class MsgTok: + """A message bound to a fixed, deterministic token rendering. + + The same MsgTok always renders to the same token segment regardless of which + turn replays it (a clean tokenizer). Token-id drift is injected explicitly by + tests via ``drift`` — never by re-rendering. + """ + + _body_counter: dict[str, int] = {} + + def __init__(self, role: str, label: str) -> None: + self.role = role + self.label = label + base = _BANDS[role] + if role == "assistant": + # An assistant message must render to the same body token as the + # response it represents (label-keyed), so a replayed assistant in a + # later prompt token-matches the original generation -> clean + # continuation. See _asst_body. + self.body = _asst_body(label) + else: + # Allocate one stable body token per (role, label). Offset past the + # END marker (base+9): the counter is shared across cases, so bodies + # must never climb into base+9 (END) or they'd collide with it. + idx = MsgTok._body_counter.setdefault(role, 0) + 1 + MsgTok._body_counter[role] = idx + self.body = base + 10 + idx + TOKEN_NAMES[self.body] = f"{role}:{_vis(label)}" + # message dict as the manager sees it (drives routing equality). + self.message = {"role": role, "content": label} + + def render(self) -> list[int]: + """[START, body, END] for this message.""" + base = _BANDS[self.role] + return [base, self.body, base + 9] + + +def sys_msg(label: str) -> MsgTok: + return MsgTok("system", label) + + +def usr_msg(label: str) -> MsgTok: + return MsgTok("user", label) + + +def asst_msg(label: str) -> MsgTok: + return MsgTok("assistant", label) + + +def tool_msg(label: str) -> MsgTok: + return MsgTok("tool", label) + + +def render_prompt(msgs: list[MsgTok]) -> list[int]: + """Render a prompt message list, appending the generation-prompt marker.""" + out: list[int] = [] + for m in msgs: + out.extend(m.render()) + out.append(_GEN) + return out + + +def render_response(label: str) -> list[int]: + """Render an assistant response: [body, ]. + + The generation-prompt marker ```` equals the assistant START token, so + `` + render_response(x)`` == the assistant message ``[, body, + ]`` replayed in a later prompt. That identity is what makes a clean + continuation hold across turns. + """ + return [_asst_body(label), _BANDS["assistant"] + 9] + + +def messages(msgs: list[MsgTok]) -> list[dict]: + """The plain message dicts record_turn wants for prompt_messages.""" + return [m.message for m in msgs] + + +def drift(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 1) -> list[int]: + """Return a copy of ``ids`` with a sentinel spliced at index ``at``. + + The sentinel sits in the drift band (7000+), so a dump shows ```` at + the exact divergence point. Splicing (insert) makes ``len`` grow by one, + which is enough to make the lcp diverge at ``at``. + """ + TOKEN_NAMES[sentinel] = "" + return ids[:at] + [sentinel] + ids[at:] + + +def drift_replace(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 2) -> list[int]: + """Return a copy of ``ids`` with the token at ``at`` REPLACED by a sentinel. + + Unlike ``drift`` this keeps length constant — used when a test wants the + divergence inside a response span without changing the cumulative length. + """ + TOKEN_NAMES[sentinel] = "" + out = list(ids) + out[at] = sentinel + return out + + +def turn(prompt_ids, response_ids, *, finish_reason="stop", logprobs=None) -> TurnRecord: + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=list(response_ids), + finish_reason=finish_reason, + output_log_probs=list(logprobs) if logprobs is not None else [], + ) + + +# A scratch space for the dual-mode printer: each case appends (title, mgr, sid, +# samples) so main() can render after the assertions pass. +_PRINT_LOG: list[tuple[str, object, str, list]] = [] + +# Raw record_turn inputs, keyed by sid, captured at call time so the printer can +# show the SOURCE data (prompt_ids / response_ids / finish / logprobs) that fed +# the tree — before any tree-building or linearization happened. +_TURN_LOG: dict[str, list[dict]] = {} + + +def _record(title: str, mgr, sid: str, samples: list) -> None: + _PRINT_LOG.append((title, mgr, sid, samples)) + + +# Convenience: append a turn with semantic messages, auto-rendering prompt unless +# an explicit prompt_ids is supplied (for drift injection). +def append( + mgr: TrajectoryManager, + sid: str, + prompt_msgs: list[MsgTok], + response_label: str | None, + *, + prompt_ids=None, + response_ids=None, + finish_reason="stop", + logprobs=None, + response_message=None, +): + p = list(prompt_ids) if prompt_ids is not None else render_prompt(prompt_msgs) + if response_ids is not None: + r = list(response_ids) + elif response_label is not None: + r = render_response(response_label) + else: + r = [] + rmsg = response_message + if rmsg is None and response_label is not None: + rmsg = {"role": "assistant", "content": response_label} + lp = logprobs + # Capture the raw turn inputs for the human-readable dump before the manager + # consumes them. + _TURN_LOG.setdefault(sid, []).append( + { + "prompt_msgs": [f"{m.role}:{_vis(m.label)}" for m in prompt_msgs], + "prompt_ids": p, + "response_ids": r, + "finish": finish_reason, + "has_lp": lp is not None, + } + ) + mgr.record_turn( + sid, + turn=turn(p, r, finish_reason=finish_reason, logprobs=lp), + prompt_messages=messages(prompt_msgs), + response_message=rmsg, + ) + return p, r + + +def _leaves(mgr, sid): + return [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + + +def _iter_all(root): + """Yield every non-root node in the tree (pre-order).""" + stack = list(root.children) + while stack: + n = stack.pop() + yield n + stack.extend(n.children) + + +# Tree text snapshot captured the instant before get_trajectory drains the sid, +# so the human-readable dump can show [tree] AND [samples] side by side even +# though get_trajectory consumes the session. +_TREE_SNAP: dict[str, str] = {} + +# Input reward passed to get_trajectory, keyed by sid, so the dump can show that +# every emitted sample carries the full input reward. +_REWARD_IN: dict[str, float] = {} + + +def get_traj(mgr, sid, *args, **kwargs): + """get_trajectory wrapper that snapshots the tree before draining. + + Linearization (get_trajectory) pops the sid, so a later dump would only see + ````. Capturing the tree text here keeps the routing tree visible + next to the Samples it produced. The input ``reward`` is captured too so the + dump can show how it maps onto the emitted samples. + """ + if mgr.has_session(sid): + _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) + _REWARD_IN[sid] = kwargs.get("reward", 0.0) + samples = mgr.get_trajectory(sid, *args, **kwargs) + # Reward assignment: get_trajectory assigns the input reward in full to every + # emitted sample (no split), so each per-sample reward must equal the input + # (modulo float error). This is the "full outcome reward per turn" invariant. + for s in samples: + assert abs(s.reward - _REWARD_IN[sid]) < 1e-9, ( + "reward not assigned in full to every sample", + s.reward, + _REWARD_IN[sid], + ) + return samples + + +def _check_invariants(samples): + for s in samples: + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length, ( + "alignment broken", + len(s.loss_mask), + len(s.rollout_log_probs), + s.response_length, + ) + assert sum(s.loss_mask) > 0, "fully-masked sample emitted" + + +def golden(sample) -> str: + """Render one Sample as a human-reviewable golden string. + + Every token is decoded to its readable name (````, ``r:done``, + ```` ...). The leading prompt prefix (no loss_mask entry) is shown as + plain names; the response region is shown with each TRAINED token (loss=1) + wrapped in ``[...]`` and each context token (loss=0) left bare. This makes the + full linearized result — tokens, where the response region starts, and + exactly which tokens carry training signal — a single literal a human can + eyeball and assert against, instead of hand-derived index arithmetic. + + Example: `` system:S user:u [r:ok] []`` + """ + toks = sample.tokens + resp_start = len(toks) - sample.response_length + parts: list[str] = [] + for i, t in enumerate(toks): + nm = name_of(t) + if i >= resp_start and sample.loss_mask[i - resp_start] == 1: + parts.append(f"[{nm}]") + else: + parts.append(nm) + return " ".join(parts) + + +def goldens(samples) -> list[str]: + return [golden(s) for s in samples] + + +# =========================================================================== +# §2 Group 1 — routing tree layer (record_turn shapes the tree) +# =========================================================================== + + +def test_1_1_single_turn_chain(): + mgr = TrajectoryManager() + sid = "1.1" + s = sys_msg("S") + u = usr_msg("compute") + p, r = append(mgr, sid, [s, u], "ok") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant"] + assert chain[-1].turn_index == 1 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:ok] []", + ] + _check_invariants(samples) + _record("1.1 single turn -> linear chain", mgr, sid, samples) + print("PASS 1.1") + + +def test_1_2_clean_multiturn_with_tool(): + mgr = TrajectoryManager() + sid = "1.2" + s, u = sys_msg("S"), usr_msg("compute") + a1, t1 = asst_msg("call"), tool_msg("4") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] + assert mgr.turn_count(sid) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:call] [] " + " tool:4 [r:done] []", + ] + _check_invariants(samples) + _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, samples) + print("PASS 1.2") + + +def test_1_3_system_fork(): + mgr = TrajectoryManager() + sid = "1.3" + for sl in ["SA", "SB"]: + append(mgr, sid, [sys_msg(sl), usr_msg("u")], "a") + root = mgr._trees[sid] + assert len(root.children) == 2, "different system -> two subtrees at root" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [ + " system:SA user:u [r:a] []", + " system:SB user:u [r:a] []", + ] + _check_invariants(samples) + _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) + print("PASS 1.3") + + +def test_1_4_user_fork_shared_system(): + mgr = TrajectoryManager() + sid = "1.4" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + root = mgr._trees[sid] + assert len(root.children) == 1, "system shared" + assert len(root.children[0].children) == 2, "user level forks" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] + _check_invariants(samples) + _record("1.4 user fork (shared system)", mgr, sid, samples) + print("PASS 1.4") + + +def test_1_5_assistant_message_fork(): + """Same (sys,user) prefix, two distinct assistant turns -> assistant fork.""" + mgr = TrajectoryManager() + sid = "1.5" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a1") + append(mgr, sid, [s, u], "a2") + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant leaves hang off shared user" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # Two independent single-turn leaves sharing only the (sys,user) prefix. + assert goldens(samples) == [ + " system:S user:u [r:a1] []", + " system:S user:u [r:a2] []", + ] + _check_invariants(samples) + _record("1.5 assistant fork under shared user", mgr, sid, samples) + print("PASS 1.5") + + +def test_1_6_tool_fork_shared_assistant(): + """Same first assistant turn, two different tool results -> tool-level fork, + making the first assistant a shared generated turn with 2 children.""" + mgr = TrajectoryManager() + sid = "1.6" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, tool_msg("x")], "ax") + append(mgr, sid, [s, u, a1, tool_msg("y")], "ay") + asst1 = mgr._trees[sid].children[0].children[0].children[0] + assert asst1.role == "assistant" and asst1.turn is not None + assert len(asst1.children) == 2, "shared assistant forks at the tool level" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # Leaf X owns the shared turn 1 (r:call trained); leaf Y shares it -> r:call + # demoted to loss=0 context, only r:ay trains (cross-leaf dedup). + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:ax] []", + " system:S user:u r:call " " tool:y [r:ay] []", + ] + _check_invariants(samples) + _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) + print("PASS 1.6") + + +def test_1_7_token_only_drift_no_fork(): + """Identical messages, tampered prompt_ids -> NO tree fork (DFS ignores + tokens), but the drift DOES surface in the linearized sample: it lands in + leaf 2's prompt region (stripped / loss=0), proving token drift cannot + corrupt a trained response yet is still carried in the sample tokens.""" + mgr = TrajectoryManager() + sid = "1.7" + s, u = sys_msg("S"), usr_msg("u") + pa, _ = append(mgr, sid, [s, u], "a") + tampered = drift(pa, 1) # spliced into the prompt at index 1 + append(mgr, sid, [s, u], "b", prompt_ids=tampered) + # Tree: (sys,user) shared, two assistant turns hang off it -> two leaves; the + # path above the assistant is single (NOT forked on tokens). + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" + assert len(mgr._trees[sid].children) == 1 + + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Leaf 1: clean. Leaf 2: the token sits in the stripped prompt region + # (bare, no brackets); the response r:b is fully trained ([...]). + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + ] + # Belt-and-suspenders on the drift placement: token present, but never inside + # the response region. + s_b = samples[1] + assert (_DRIFT_BAND + 1) in s_b.tokens, "drift token is still carried in the sample" + assert (_DRIFT_BAND + 1) not in s_b.tokens[len(tampered) :], "drift not in the response region" + _check_invariants(samples) + _record("1.7 token-only drift -> no tree fork, drift lands in stripped prompt", mgr, sid, samples) + print("PASS 1.7") + + +def test_1_8_multi_tool_per_turn(): + mgr = TrajectoryManager() + sid = "1.8" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + ta, tb = tool_msg("A"), tool_msg("B") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, ta, tb], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "tool", "assistant"] + assert chain[3].message == ta.message + assert chain[4].message == tb.message + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:A tool:B [r:done] []", + ] + _check_invariants(samples) + _record("1.8 multi-tool turn -> one node per tool", mgr, sid, samples) + print("PASS 1.8") + + +def test_1_9_cross_sid_isolation(): + mgr = TrajectoryManager() + s = sys_msg("S") + for sid, ul in [("sid-a", "A"), ("sid-b", "B")]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + assert len(_leaves(mgr, "sid-a")) == 1 + assert len(_leaves(mgr, "sid-b")) == 1 + assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] + sa = get_traj(mgr, "sid-a", base_sample=Sample(index=0, prompt=""), reward=1.0) + sb = get_traj(mgr, "sid-b", base_sample=Sample(index=1, prompt=""), reward=1.0) + assert goldens(sa) == [" system:S user:A [r:a] []"] + assert goldens(sb) == [" system:S user:B [r:b] []"] + _check_invariants(sa) + _check_invariants(sb) + _record("1.9 cross-sid isolation (sid-a)", mgr, "sid-a", sa) + _record("1.9 cross-sid isolation (sid-b)", mgr, "sid-b", sb) + print("PASS 1.9") + + +def test_1_10_empty_response(): + mgr = TrajectoryManager() + sid = "1.10" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], None, response_ids=[], response_message=None, finish_reason="length") + asst = _leaves(mgr, sid)[0] + assert asst.role == "assistant" + assert asst.turn.output_ids == [] + assert asst.message is None + # Empty response -> the only turn has no trainable token, so its segment is + # dropped at linearization (no fully-masked sample). Zero samples is correct. + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 0 + _record("1.10 empty response -> assistant leaf, no message (0 samples)", mgr, sid, samples) + print("PASS 1.10") + + +# =========================================================================== +# §2 Group 2 — linearization layer (get_trajectory token routing) +# =========================================================================== + + +def test_2_1_single_turn_linearize(): + mgr = TrajectoryManager() + sid = "2.1" + s, u = sys_msg("S"), usr_msg("u") + p, r = append(mgr, sid, [s, u], "a", logprobs=None) + # attach explicit logprobs so we can check propagation + leaf = _leaves(mgr, sid)[0] + leaf.turn = dataclasses.replace(leaf.turn, output_log_probs=[-0.5] * len(r)) + samples = get_traj(mgr, sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + assert goldens(samples) == [" system:S user:u [r:a] []"] + assert s0.rollout_log_probs == [-0.5] * len(r) + assert s0.reward == 1.0 + _check_invariants(samples) + _record("2.1 single-turn linearize", mgr, sid, samples) + print("PASS 2.1") + + +def test_2_2_clean_multiturn_linearize(): + mgr = TrajectoryManager() + sid = "2.2" + s, u, a1, t1 = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("4") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _common_prefix_len(p1 + r1, p2) + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:4 [r:done] []", + ] + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.2 clean 2-turn linearize", mgr, sid, samples) + print("PASS 2.2") + + +def test_2_3_drift_case_A_forks(): + """Drift inside a PROMPT region -> case A -> fork, no token dropped.""" + mgr = TrajectoryManager() + sid = "2.3" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # case-A fork: two coherent single-turn segments; the token stays in + # segment 2's stripped prompt region (bare), no token dropped. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) + print("PASS 2.3") + + +def test_2_4_drift_case_B1_short_replaces(): + """Small drift inside the most-recent response span -> replace.""" + mgr = TrajectoryManager() # default threshold 1024 + sid = "2.4" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + assert p2_honest[: len(p1) + len(r1)] == p1 + r1 + drift_idx = len(p1) + len(r1) - 1 # last token of r1's echo + p2 = drift_replace(p2_honest, drift_idx) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _common_prefix_len(p1 + r1, p2) + assert L == drift_idx + # replace: the drifted r:call response is no longer a faithful echo of what the + # model generated (its tail diverged), so the WHOLE surviving span is masked and + # re-supplied as loss=0 prompt context (the token marks the divergence); + # only the new r:done trains. + assert goldens(samples) == [ + " system:S user:u r:call " + " tool:t [r:done] []", + ] + assert s0.rollout_log_probs == [0.0] * (len(p2) - len(p1)) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.4 drift case B1 (small) -> replace", mgr, sid, samples) + print("PASS 2.4") + + +def test_2_5_drift_case_B1_long_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=1) # d>=1 -> fork + sid = "2.5" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Both segments single-turn (the drift forked them apart): each trains its own + # response. The sits in segment 2's stripped prompt. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) + print("PASS 2.5") + + +def test_2_6_drift_case_B1_threshold_zero_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=0) + sid = "2.6" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) + print("PASS 2.6") + + +def test_2_7_drift_case_B2_earlier_turn_forks(): + """Drift inside an EARLIER turn's response span -> always fork.""" + mgr = TrajectoryManager() + sid = "2.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2, t2 = asst_msg("a2"), tool_msg("t2") + p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + p3_honest = render_prompt([s, u, a1, t1, a2, t2]) + p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) + p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Segment 1 = clean turns 1+2; segment 2 = turn 3 alone (forked because the + # drift hit an EARLIER turn's response span, which replace can't drop). + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:t2 [r:a3] []", + ] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) + print("PASS 2.7") + + +def test_2_8_fork_reward_split(): + mgr = TrajectoryManager() + sid = "2.8" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # case-A fork: two single-turn segments, each trains its own response. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + # reward 1.0 assigned in full to each of the 2 forked samples. + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.8 fork reward (1.0 to each sample)", mgr, sid, samples) + print("PASS 2.8") + + +def test_2_9_two_leaves_reward_split(): + mgr = TrajectoryManager() + sid = "2.9" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] + # reward 1.0 assigned in full to each of the 2 leaves. + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.9 two leaves reward (1.0 to each sample)", mgr, sid, samples) + print("PASS 2.9") + + +def test_2_10_cross_leaf_dedup(): + """Shared assistant trained on first leaf only; second leaf re-emits it + as loss=0 context.""" + mgr = TrajectoryManager() + sid = "2.10" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + tx, ty = tool_msg("x"), tool_msg("y") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) + p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + s_second = samples[1] + # First leaf trains the shared r:call + its own r:a2; second leaf shares + # r:call (demoted to loss=0) and trains only r:a3. + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:a2] []", + " system:S user:u r:call " " tool:y [r:a3] []", + ] + assert s_second.rollout_log_probs == [0.0] * (len(p3) - len(p1)) + [-0.3] * len(r3) + _check_invariants(samples) + _record("2.10 cross-leaf dedup (shared assistant trained once)", mgr, sid, samples) + print("PASS 2.10") + + +def test_2_11_routing_only_assistant_filtered(): + """cc replays an assistant the manager never recorded -> mounts routing-only, + must be filtered out of the strict-prefix walk (no raise).""" + mgr = TrajectoryManager() + sid = "2.11" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls") + foreign = asst_msg("foreign") + t2 = tool_msg("t2") + append(mgr, sid, [s, u, a1, t1, a2, foreign, t2], "a3") + leaves = _leaves(mgr, sid) + assert len(leaves) == 1 + chain = leaves[0].path_from_root() + routing = [n for n in chain if n.role == "assistant" and n.turn is None] + assert len(routing) == 1 and routing[0].message["content"] == "foreign" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # The foreign assistant (r:foreign) is routing-only -> appears as bare context + # (no brackets); the three real turns r:a1/r:a2/r:a3 train. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] r:foreign " + " tool:t2 [r:a3] []", + ] + _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) + print("PASS 2.11") + + +def test_2_12_drop_clears_sid(): + mgr = TrajectoryManager() + sid = "2.12" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + assert mgr.has_session(sid) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [" system:S user:u [r:a] []"] + assert not mgr.has_session(sid) + assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] + _check_invariants(samples) + _record("2.12 drop clears sid (2nd get_trajectory -> [])", mgr, sid, samples) + print("PASS 2.12") + + +# =========================================================================== +# §2 Group 3 — combined / stress (both layers interacting) +# =========================================================================== + + +def test_3_1_rewrite_merge_absorbs_short(): + mgr = TrajectoryManager() + sid = "3.1" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok ") # cc-rewritten (different message identity) + t1 = tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "short rewrite absorbed, not forked" + chain = leaves[0].path_from_root() + merged = chain[2] + assert merged.turn is None and merged.turn_index is None + assert merged.message == a1_rw.message + assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # The abandoned turn-1 response (r:ok␣) is demoted to routing-only -> appears + # bare; only the surviving turn-2 r:done trains. + assert goldens(samples) == [ + " system:S user:u r:ok␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.1 rewrite-merge absorbs short assistant", mgr, sid, samples) + print("PASS 3.1") + + +def test_3_2_rewrite_merge_long_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=1) # r1 len 2 >= 1 + sid = "3.2" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok2 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Leaf 1: the abandoned turn-1 standalone (r:ok). Leaf 2: turn-2 only (the + # rewritten r:ok2␣ assistant mounts routing-only and is filtered out). + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok2␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) + print("PASS 3.2") + + +def test_3_3_rewrite_merge_threshold_zero_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=0) + sid = "3.3" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok3 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok3␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) + print("PASS 3.3") + + +def test_3_4_rewrite_merge_ambiguous_forks(): + mgr = TrajectoryManager() + sid = "3.4" + s, u = sys_msg("S"), usr_msg("u") + # two short assistant leaves under shared (sys,user) + append(mgr, sid, [s, u], "a") + append(mgr, sid, [s, u], "b") + a_c, t1 = asst_msg("c"), tool_msg("t") + append(mgr, sid, [s, u, a_c, t1], "d") + assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3 + # Leaves "a" and "b" are standalone single turns; leaf "d" carries the + # ambiguous-rewrite assistant (r:c) as routing-only (bare) -> trains only r:d. + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + " system:S user:u r:c " " tool:t [r:d] []", + ] + _check_invariants(samples) + _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) + print("PASS 3.4") + + +def test_3_5_rewrite_merge_match_key_updated(): + """After merge, a later turn replaying the rewritten message must descend + through the merged node (match_key updated), not fork again.""" + mgr = TrajectoryManager() + sid = "3.5" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok5 ") + t1, a2, t2 = tool_msg("t1"), asst_msg("second"), tool_msg("t2") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "match_key updated -> no spurious fork" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # Turn 1 (r:ok) was absorbed as routing-only (rewrite merge), so it appears + # bare; turns 2 and 3 (r:second / r:third) train in one clean chain. + assert goldens(samples) == [ + " system:S user:u r:ok5␣ " + " tool:t1 [r:second] [] tool:t2 [r:third] []", + ] + _check_invariants(samples) + _record("3.5 rewrite-merge match_key updated", mgr, sid, samples) + print("PASS 3.5") + + +def test_3_6_tree_fork_plus_token_drift(): + """A tree fork (two leaves) where ONE leaf also drift-forks internally, + yielding 3 Samples total. Combines layer-1 (message fork) with layer-2 + (token drift fork).""" + mgr = TrajectoryManager() + sid = "3.6" + s, u = sys_msg("S"), usr_msg("u") + a1, tx, ty = asst_msg("call"), tool_msg("x"), tool_msg("y") + ax2 = asst_msg("ax2") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + # Leaf X: clean continuation, then a third turn with a case-A prompt drift. + append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + txx = tool_msg("xx") + p3_honest = render_prompt([s, u, a1, tx, ax2, txx]) + p3 = drift(p3_honest, len(p1) - 1) # case A drift -> fork inside leaf X + append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) + # Leaf Y: a separate tool result off the shared assistant. + append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3, [s.tokens for s in samples] + assert goldens(samples) == [ + # Sample 0: leaf X first segment, trains the shared r:call (first claim) + r:ax2. + " system:S user:u [r:call] [] " + " tool:x [r:ax2] []", + # Sample 1: leaf X second segment, a FRESH segment after the case-A fork -> + # whole prompt stripped, only r:ax3 trains (the sits in its prompt). + " system:S user:u r:call " + " tool:x r:ax2 tool:xx [r:ax3] []", + # Sample 2: leaf Y, shares r:call (claimed by sample 0 -> bare), trains r:ay2. + " system:S user:u r:call " " tool:y [r:ay2] []", + ] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) + print("PASS 3.6") + + +def test_3_7_deep_multi_leaf_dedup(): + """Three leaves sharing a 2-level assistant prefix; the shared turns are + trained exactly once across all leaves.""" + mgr = TrajectoryManager() + sid = "3.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + # three different tool results off a2 -> three leaves sharing a1+a2 + for lbl in ["p", "q", "r"]: + append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3 + # Leaf 0 OWNS the shared r:a1 + r:a2 (both trained) and its own end-p; leaves + # 1 and 2 SHARE r:a1 + r:a2 (bare, claimed by leaf 0) and train only their own + # end response. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] tool:p [r:end-p] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:q [r:end-q] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:r [r:end-r] []", + ] + _check_invariants(samples) + _record("3.7 deep multi-leaf dedup (3 leaves, shared trained once)", mgr, sid, samples) + print("PASS 3.7") + + +def test_3_8_long_mixed_session(): + """A ~7-turn session combining clean continuation, a mid-session B1 replace, + and a final case-A fork — verifying the mechanisms chain without interfering.""" + mgr = TrajectoryManager() + sid = "3.8" + s, u = sys_msg("S"), usr_msg("u") + a = [asst_msg(f"a{i}") for i in range(6)] + t = [tool_msg(f"t{i}") for i in range(6)] + lp = [-0.5, -0.5] + # turn 1 + p1, r1 = append(mgr, sid, [s, u], "a0", finish_reason="tool_calls", logprobs=lp) + # turns 2..4 clean + prefix = [s, u, a[0], t[0]] + append(mgr, sid, prefix, "a1", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[1], t[1]] + append(mgr, sid, prefix, "a2", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[2], t[2]] + # turn 5: B1 small replace — drift the last token of the previous response. + p5_honest = render_prompt(prefix) + p5 = drift_replace(p5_honest, len(p5_honest) - 2) # near tail, inside last resp echo region + append(mgr, sid, prefix, "a3", prompt_ids=p5, finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[3], t[3]] + # turn 6: clean + append(mgr, sid, prefix, "a4", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[4], t[4]] + # turn 7: case-A fork (drift in early prompt region) + p7_honest = render_prompt(prefix) + p7 = drift(p7_honest, len(p1) - 1) + append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # The final case-A fork splits the single leaf chain into 3 segments. + assert goldens(samples) == [ + # Segment 1: turns 1-4 in one clean chain. The turn-5 B1 replace dropped a + # drifted tail (the is gone here) and realigned, so r:a0..r:a3 all + # train. + " system:S user:u [r:a0] [] " + " tool:t0 [r:a1] [] tool:t1 [r:a2] [] " + " tool:t2 [r:a3] []", + # Segment 2: cross-leaf-style dedup within the chain — the prior turns are + # re-emitted as bare context and only r:a4 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 [r:a4] []", + # Segment 3: turn 7 after the case-A fork (the in the early prompt + # region); whole prefix bare, only r:a5 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 r:a4 " + " tool:t4 [r:a5] []", + ] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) + print("PASS 3.8") + + +# =========================================================================== +# §2 Group 4 — boundary / defensive / feature-completion +# +# Fills coverage gaps the matrix above left open: +# input-validation contracts, mixed-logprobs trajectories, the case-B1 drift +# threshold boundary, and the default-base_sample path. +# =========================================================================== + + +def test_4_2_logprobs_length_mismatch_raises(): + """output_log_probs whose length != output_ids -> ValueError at record_turn.""" + mgr = TrajectoryManager() + sid = "4.2" + s, u = sys_msg("S"), usr_msg("u") + bad = TurnRecord( + prompt_ids=render_prompt([s, u]), + output_ids=[9101, 9102, 9103], + finish_reason="stop", + output_log_probs=[-0.1, -0.2], # length 2 != 3 + ) + raised = False + try: + mgr.record_turn( + sid, + turn=bad, + prompt_messages=messages([s, u]), + response_message={"role": "assistant", "content": "x"}, + ) + except AssertionError as e: + raised = True + assert "output_log_probs" in str(e) + assert raised, "expected AssertionError on logprobs/ids length mismatch" + print("PASS 4.2") + + +def test_4_3_empty_prompt_messages_skipped(): + """Empty prompt_messages -> record_turn is a no-op (warns, no node, no turn).""" + mgr = TrajectoryManager() + sid = "4.3" + mgr.record_turn( + sid, + turn=turn([1], [2], finish_reason="stop"), + prompt_messages=[], + response_message=None, + ) + assert mgr.turn_count(sid) == 0 + # The tree may be created empty (root only) or absent; either way no leaf. + assert not mgr.has_session(sid) or list(_leaves(mgr, sid)) == [] + print("PASS 4.3") + + +def test_4_4_default_base_sample(): + """get_trajectory without base_sample is rejected (required keyword-only arg).""" + mgr = TrajectoryManager() + sid = "4.4" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + with pytest.raises(TypeError): + mgr.get_trajectory(sid, reward=1.0) # no base_sample + print("PASS 4.4") + + +def test_4_5_mixed_logprobs_across_turns(): + """A trajectory where turn 1 carries logprobs and turn 2 does NOT: the + sample's turn-1 response region has real logprobs, the turn-2 region is + padded with 0.0 (the response is still trained, loss=1).""" + mgr = TrajectoryManager() + sid = "4.5" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=None) # no logprobs + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _common_prefix_len(p1 + r1, p2) + # both responses still trained (golden shows the loss layout)... + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:t [r:done] []", + ] + # ...but turn-2's region carries padded 0.0 logprobs (it had none), while + # turn-1's region keeps its real logprobs. + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [0.0] * len(r2) + _check_invariants(samples) + _record("4.5 mixed logprobs across turns (turn2 padded 0.0)", mgr, sid, samples) + print("PASS 4.5") + + +def test_4_6_drift_B1_threshold_boundary(): + """case-B1 threshold compares the incoming turn's full ``output_ids`` length + to ``fork_threshold`` (mirroring ``_try_merge_assistant_rewrite``): the gate + is exclusive, so ``len(r2) == threshold`` forks and ``len(r2) < threshold`` + replaces. Drift-tail length is not part of the gate -- only its position + (inside the most-recent response span) keeps REALIGN physically applicable.""" + + def run(threshold, new_resp_len): + mgr = TrajectoryManager(fork_threshold_tokens=threshold) + sid = f"4.6-{threshold}-{new_resp_len}" + s, u = sys_msg("S"), usr_msg("u") + # 4-token response so the divergence can sit inside the response span. + p1 = render_prompt([s, u]) + r1 = [9001, 9002, 9003, 9004] + mgr.record_turn( + sid, + turn=turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=messages([s, u]), + response_message={"role": "assistant", "content": "a1"}, + ) + a1m = {"role": "assistant", "content": "a1"} + tm = tool_msg("t") + # honest turn-2 prompt echoes p1 + r1 then the tool block + gen marker. + p2_honest = p1 + r1 + tm.render() + [_GEN] + # divergence one token before the end of r1's echo (well inside the response span). + drift_idx = len(p1) + len(r1) - 1 + p2 = drift_replace(p2_honest, drift_idx) + r2 = [9100 + i for i in range(new_resp_len)] + mgr.record_turn( + sid, + turn=turn(p2, r2, finish_reason="stop"), + prompt_messages=[*messages([s, u]), a1m, tm.message], + response_message={"role": "assistant", "content": "done"}, + ) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + return samples, p1, r1, p2, r2 + + # len(r2) == threshold -> fork: two single-turn segments, each trains its own resp. + forked, p1, r1, p2, r2 = run(threshold=2, new_resp_len=2) + assert len(forked) == 2, f"len(r2)==threshold must fork, got {len(forked)}" + assert forked[0].tokens == p1 + r1 + assert forked[0].loss_mask == [1] * len(r1) + assert forked[1].tokens == p2 + r2 + assert forked[1].loss_mask == [1] * len(r2) + # len(r2) < threshold -> replace: one coherent segment realigned to p2. + replaced, p1b, r1b, p2b, r2b = run(threshold=3, new_resp_len=2) + assert len(replaced) == 1, f"len(r2) the WHOLE r1 span is + # masked (loss=0 prompt context), r2 trains. + assert replaced[0].loss_mask == [0] * (len(p2b) - len(p1b)) + [1] * len(r2b) + _check_invariants(forked) + _check_invariants(replaced) + print("PASS 4.6") + + +# =========================================================================== +# §3 Dual-mode printer +# =========================================================================== + + +def _print_sample(idx: int, s: Sample) -> None: + toks = s.tokens + resp_start = len(toks) - s.response_length + # build aligned token/loss rows over the response region (the trained part); + # the leading prompt prefix has no loss_mask entry. + names = [name_of(t) for t in toks] + loss = ["-"] * resp_start + [str(x) for x in s.loss_mask] + widths = [max(len(names[i]), len(loss[i])) for i in range(len(toks))] + tok_row = " ".join(names[i].ljust(widths[i]) for i in range(len(toks))) + loss_row = " ".join(loss[i].ljust(widths[i]) for i in range(len(toks))) + print(f" Sample#{idx} reward={s.reward:.3f} resp_len={s.response_length}") + print(f" tok : {tok_row}") + print(f" loss: {loss_row}") + + +def _print_raw_turns(sid: str) -> None: + """Print the raw record_turn inputs (the SOURCE data) for a sid. + + Shows, per turn, the prompt message labels and the actual prompt_ids / + response_ids decoded to readable names, plus finish_reason and whether + logprobs were attached. This is what fed the tree, before any building or + linearization. + """ + turns = _TURN_LOG.get(sid, []) + print(f"[raw turns] {len(turns)}") + for k, t in enumerate(turns, start=1): + msgs = " , ".join(t["prompt_msgs"]) + print(f" turn#{k} finish={t['finish']} has_logprobs={t['has_lp']}") + print(f" msgs : {msgs}") + print(f" prompt : {render_ids(t['prompt_ids'])}") + print(f" output : {render_ids(t['response_ids']) or ''}") + + +def _print_case(title: str, mgr, sid: str, samples: list) -> None: + print(f"\n=== CASE {title} ===") + _print_raw_turns(sid) + if mgr.has_session(sid): + txt = dump_tree_txt(mgr, sid) + else: + # Session already drained by get_trajectory; fall back to the snapshot + # captured by get_traj just before draining. + txt = _TREE_SNAP.get(sid, "") + print("[tree]") + for line in txt.splitlines(): + print(" " + line) + n = len(samples) + if n: + r_in = _REWARD_IN.get(sid, 0.0) + print(f"[samples] {n} (reward: {r_in:.3f} assigned in full to each sample)") + else: + print(f"[samples] {n}") + for i, s in enumerate(samples): + _print_sample(i, s) + + +# =========================================================================== +# main +# =========================================================================== + + +_CASES = [ + test_1_1_single_turn_chain, + test_1_2_clean_multiturn_with_tool, + test_1_3_system_fork, + test_1_4_user_fork_shared_system, + test_1_5_assistant_message_fork, + test_1_6_tool_fork_shared_assistant, + test_1_7_token_only_drift_no_fork, + test_1_8_multi_tool_per_turn, + test_1_9_cross_sid_isolation, + test_1_10_empty_response, + test_2_1_single_turn_linearize, + test_2_2_clean_multiturn_linearize, + test_2_3_drift_case_A_forks, + test_2_4_drift_case_B1_short_replaces, + test_2_5_drift_case_B1_long_forks, + test_2_6_drift_case_B1_threshold_zero_forks, + test_2_7_drift_case_B2_earlier_turn_forks, + test_2_8_fork_reward_split, + test_2_9_two_leaves_reward_split, + test_2_10_cross_leaf_dedup, + test_2_11_routing_only_assistant_filtered, + test_2_12_drop_clears_sid, + test_3_1_rewrite_merge_absorbs_short, + test_3_2_rewrite_merge_long_forks, + test_3_3_rewrite_merge_threshold_zero_forks, + test_3_4_rewrite_merge_ambiguous_forks, + test_3_5_rewrite_merge_match_key_updated, + test_3_6_tree_fork_plus_token_drift, + test_3_7_deep_multi_leaf_dedup, + test_3_8_long_mixed_session, + test_4_2_logprobs_length_mismatch_raises, + test_4_3_empty_prompt_messages_skipped, + test_4_4_default_base_sample, + test_4_5_mixed_logprobs_across_turns, + test_4_6_drift_B1_threshold_boundary, +] + + +def main() -> None: + for case in _CASES: + case() + # Replay the captured tree / sample snapshots as human-readable dumps. + print("\n" + "=" * 70) + print("HUMAN-READABLE DUMPS") + print("=" * 70) + for title, mgr, sid, samples in _PRINT_LOG: + _print_case(title, mgr, sid, samples) + print(f"\nALL CASES PASSED ({len(_CASES)} cases)") + + +if __name__ == "__main__": + if os.environ.get("TRAJ_DUMP"): + main() # human-readable tree / sample dumps + else: + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_agent_adapters.py b/tests/test_agent_adapters.py deleted file mode 100644 index c234b4c32..000000000 --- a/tests/test_agent_adapters.py +++ /dev/null @@ -1,853 +0,0 @@ -import asyncio -import json -import sys -from pathlib import Path - -import pytest -from aiohttp import web -from aiohttp.test_utils import TestClient, TestServer - -REPO_ROOT = Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from vime.agent.adapters import anthropic, openai -from vime.agent.adapters.common import VLLM_URL_KEY -from vime.agent.trajectory import TurnRecord - - -NUM_GPUS = 0 - - -class ToyTokenizer: - def __init__(self, outputs: dict[tuple[int, ...], str] | None = None) -> None: - self.outputs = outputs or {} - self.rendered: list[tuple[list[dict], list[dict] | None]] = [] - - def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): - self.rendered.append((list(messages), tools)) - return list(range(1, len(messages) + 2)) - - def decode(self, ids, skip_special_tokens=False): - return self.outputs.get(tuple(ids), "") - - -class ScriptedTokenizer(ToyTokenizer): - def __init__(self, prompts: list[list[int]], outputs: dict[tuple[int, ...], str]) -> None: - super().__init__(outputs) - self.prompts = [list(prompt) for prompt in prompts] - - def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): - self.rendered.append((list(messages), tools)) - assert self.prompts, "unexpected chat-template render" - return self.prompts.pop(0) - - -class FakeVLLM: - def __init__(self, turns: list[list[tuple[float, int]]]) -> None: - self.turns = [list(turn) for turn in turns] - self.requests: list[dict] = [] - self.routing_keys: list[str | None] = [] - - async def handle_generate(self, request): - self.routing_keys.append(request.headers.get("x-session-id")) - self.requests.append(await request.json()) - assert self.turns, "unexpected /inference/v1/generate call" - turn = self.turns.pop(0) - token_ids = [token_id for _logprob, token_id in turn] - content = [{"logprob": logprob} for logprob, _token_id in turn] - return web.json_response( - { - "choices": [ - { - "token_ids": token_ids, - "logprobs": {"content": content}, - "finish_reason": "stop", - } - ] - } - ) - - -class FakeRequest: - def __init__(self, headers: dict[str, str]) -> None: - self.headers = headers - - -def _parse_sse(raw: str) -> list[tuple[str, object]]: - events: list[tuple[str, object]] = [] - event_name = "message" - data_lines: list[str] = [] - - def flush() -> None: - nonlocal event_name, data_lines - if not data_lines: - event_name = "message" - return - data = "\n".join(data_lines) - payload: object - if data == "[DONE]": - payload = data - else: - payload = json.loads(data) - events.append((event_name, payload)) - event_name = "message" - data_lines = [] - - for line in raw.splitlines(): - if not line: - flush() - elif line.startswith("event:"): - event_name = line.removeprefix("event:").strip() - elif line.startswith("data:"): - data_lines.append(line.removeprefix("data:").strip()) - flush() - return events - - -@pytest.mark.unit -def test_session_id_comes_from_protocol_fields_not_custom_header(): - assert ( - openai._request_session_id( - FakeRequest({"X-Vime-Session-Id": "custom"}), - {"metadata": {"session_id": "meta-session"}, "user": "body-user"}, - ) - == "meta-session" - ) - assert ( - openai._request_session_id(FakeRequest({"X-Vime-Session-Id": "custom"}), {"user": "body-user"}) == "body-user" - ) - assert ( - anthropic._request_session_id(FakeRequest({"X-Vime-Session-Id": "custom", "X-Api-Key": "anthropic-key"})) - == "anthropic-key" - ) - assert ( - anthropic._request_session_id( - FakeRequest({"Authorization": "Bearer bearer-session", "X-Api-Key": "anthropic-key"}) - ) - == "bearer-session" - ) - - -@pytest.mark.unit -def test_anthropic_translation_keeps_tool_results_and_tool_schema(): - messages = [ - {"role": "user", "content": [{"type": "text", "text": "hi"}]}, - { - "role": "assistant", - "content": [ - {"type": "thinking", "thinking": "plan"}, - {"type": "text", "text": "ok"}, - {"type": "tool_use", "name": "lookup", "input": {"q": "vime"}}, - ], - }, - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "u1", "content": "result"}]}, - ] - - translated = anthropic._translate_anthropic(messages, system="sys") - tools = anthropic._anthropic_tools_to_chat_tools( - [{"name": "lookup", "description": "search", "input_schema": {"type": "object"}}] - ) - - assert translated == [ - {"role": "system", "content": "sys"}, - {"role": "user", "content": "hi"}, - { - "role": "assistant", - "content": "ok", - "reasoning_content": "plan", - "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "vime"}}}], - }, - {"role": "tool", "content": "result"}, - ] - assert tools == [ - { - "type": "function", - "function": { - "name": "lookup", - "description": "search", - "parameters": {"type": "object"}, - }, - } - ] - - -@pytest.mark.unit -def test_openai_translation_and_responses_input_shapes(): - chat_messages = openai._translate_chat_messages( - [ - {"role": "developer", "content": "rules"}, - {"role": "user", "content": [{"type": "text", "text": "hello"}]}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "lookup", "arguments": {"q": "vime"}}, - } - ], - }, - {"role": "tool", "tool_call_id": "call_1", "content": "found"}, - ] - ) - response_messages = openai._responses_input_to_messages( - [ - {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, - {"type": "function_call_output", "call_id": "call_1", "output": "answer"}, - ], - instructions="be brief", - ) - - assert chat_messages == [ - {"role": "system", "content": "rules"}, - {"role": "user", "content": "hello"}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "lookup", "arguments": '{"q": "vime"}'}, - } - ], - }, - {"role": "tool", "content": "found", "tool_call_id": "call_1"}, - ] - assert response_messages == [ - {"role": "system", "content": "be brief"}, - {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, - {"role": "tool", "tool_call_id": "call_1", "content": "answer"}, - ] - - -@pytest.mark.unit -def test_openai_chat_completion_endpoint_records_token_segments(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[101], finish_reason="stop", output_log_probs=[-0.1]) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = ToyTokenizer({(101,): "hello"}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-chat", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/chat/completions", - headers={"Authorization": "Bearer sid-chat"}, - json={ - "model": "actor", - "messages": [{"role": "user", "content": "hello?"}], - "max_tokens": 4, - }, - ) - data = await resp.json() - finally: - await client.close() - - segments = await adapter.finish_session("sid-chat") - assert resp.status == 200 - assert data["object"] == "chat.completion" - assert data["choices"][0]["message"] == {"role": "assistant", "content": "hello"} - assert data["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} - assert segments[0].prompt_ids == [1, 2] - assert segments[0].response_ids == [101] - assert segments[0].loss_mask == [1] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_chat_completion_streaming_returns_sse_chunks_and_records_segments(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[401], finish_reason="stop", output_log_probs=[-0.4]) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = ToyTokenizer({(401,): "streamed text"}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-chat-stream", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/chat/completions", - headers={"Authorization": "Bearer sid-chat-stream"}, - json={ - "model": "actor", - "stream": True, - "messages": [{"role": "user", "content": "hello?"}], - }, - ) - raw = await resp.text() - finally: - await client.close() - - events = _parse_sse(raw) - chunks = [payload for _, payload in events if isinstance(payload, dict)] - segments = await adapter.finish_session("sid-chat-stream") - assert resp.status == 200 - assert chunks[0]["object"] == "chat.completion.chunk" - assert chunks[0]["choices"][0]["delta"] == {"role": "assistant"} - assert any(c["choices"][0]["delta"] == {"content": "streamed text"} for c in chunks) - assert chunks[-1]["choices"][0]["finish_reason"] == "stop" - assert chunks[-1]["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} - assert events[-1] == ("message", "[DONE]") - assert segments[0].prompt_ids == [1, 2] - assert segments[0].response_ids == [401] - assert segments[0].rollout_log_probs == [-0.4] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_chat_completion_streaming_returns_tool_call_delta(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[451], finish_reason="stop", output_log_probs=[-0.45] - ) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "use it vime" - tokenizer = ToyTokenizer({(451,): raw}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-chat-tool-stream", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/chat/completions", - headers={"Authorization": "Bearer sid-chat-tool-stream"}, - json={ - "model": "actor", - "stream": True, - "messages": [{"role": "user", "content": "call lookup"}], - "tools": [ - { - "type": "function", - "function": { - "name": "lookup", - "description": "search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - }, - } - ], - }, - ) - raw_sse = await resp.text() - finally: - await client.close() - - chunks = [payload for _, payload in _parse_sse(raw_sse) if isinstance(payload, dict)] - tool_delta = next(c["choices"][0]["delta"] for c in chunks if "tool_calls" in c["choices"][0]["delta"]) - segments = await adapter.finish_session("sid-chat-tool-stream") - assert resp.status == 200 - assert any(c["choices"][0]["delta"] == {"content": "use it"} for c in chunks) - assert tool_delta["tool_calls"][0]["index"] == 0 - assert tool_delta["tool_calls"][0]["function"]["name"] == "lookup" - assert tool_delta["tool_calls"][0]["function"]["arguments"] == '{"query": "vime"}' - assert chunks[-1]["choices"][0]["finish_reason"] == "tool_calls" - assert segments[0].response_ids == [451] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_responses_endpoint_returns_function_calls(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[301], finish_reason="stop", output_log_probs=[-0.3]) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "look vime" - tokenizer = ToyTokenizer({(301,): raw}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-responses", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/responses", - headers={"Authorization": "Bearer sid-responses"}, - json={ - "model": "actor", - "input": "find it", - "tools": [ - { - "type": "function", - "name": "lookup", - "description": "search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - data = await resp.json() - finally: - await client.close() - - output_types = [item["type"] for item in data["output"]] - function_call = next(item for item in data["output"] if item["type"] == "function_call") - segments = await adapter.finish_session("sid-responses") - assert resp.status == 200 - assert data["object"] == "response" - assert output_types == ["message", "function_call"] - assert data["output"][0]["content"][0]["text"] == "look" - assert function_call["name"] == "lookup" - assert function_call["arguments"] == '{"query": "vime"}' - assert segments[0].response_ids == [301] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_responses_streaming_preserves_function_call_output(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[551], finish_reason="stop", output_log_probs=[-0.55] - ) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "vime" - tokenizer = ToyTokenizer({(551,): raw}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-responses-tool-stream", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/responses", - headers={"Authorization": "Bearer sid-responses-tool-stream"}, - json={ - "model": "actor", - "stream": True, - "input": "call lookup", - "tools": [ - { - "type": "function", - "name": "lookup", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - raw_sse = await resp.text() - finally: - await client.close() - - events = _parse_sse(raw_sse) - created = next(payload for name, payload in events if name == "response.created") - completed = next(payload for name, payload in events if name == "response.completed") - completed_call = next(item for item in completed["response"]["output"] if item["type"] == "function_call") - segments = await adapter.finish_session("sid-responses-tool-stream") - assert resp.status == 200 - assert created["type"] == "response.created" - assert [item["type"] for item in created["response"]["output"]] == ["function_call"] - assert completed_call["name"] == "lookup" - assert completed_call["arguments"] == '{"query": "vime"}' - assert segments[0].response_ids == [551] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_responses_streaming_returns_sse_events_and_records_segments(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[501], finish_reason="stop", output_log_probs=[-0.5]) - - async def run_case(): - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = ToyTokenizer({(501,): "response text"}) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-responses-stream", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/responses", - headers={"Authorization": "Bearer sid-responses-stream"}, - json={ - "model": "actor", - "stream": True, - "instructions": "be brief", - "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello?"}]}], - }, - ) - raw = await resp.text() - finally: - await client.close() - - events = _parse_sse(raw) - event_names = [name for name, _ in events] - text_delta = next(payload for name, payload in events if name == "response.output_text.delta") - completed = next(payload for name, payload in events if name == "response.completed") - segments = await adapter.finish_session("sid-responses-stream") - assert resp.status == 200 - assert event_names == ["response.created", "response.output_text.delta", "response.completed"] - assert text_delta == {"type": "response.output_text.delta", "delta": "response text"} - assert completed["response"]["status"] == "completed" - assert completed["response"]["usage"] == {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4} - assert segments[0].prompt_ids == [1, 2, 3] - assert segments[0].response_ids == [501] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_anthropic_messages_endpoint_returns_non_stream_json_and_records_segments(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[581], finish_reason="stop", output_log_probs=[-0.58] - ) - - async def run_case(): - monkeypatch.setattr(anthropic, "_generate", fake_generate) - tokenizer = ToyTokenizer({(581,): "plain response"}) - adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-anthropic-json", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/messages", - headers={"Authorization": "Bearer sid-anthropic-json"}, - json={ - "model": "actor", - "system": "be useful", - "max_tokens": 4, - "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], - }, - ) - data = await resp.json() - finally: - await client.close() - - segments = await adapter.finish_session("sid-anthropic-json") - assert resp.status == 200 - assert data["type"] == "message" - assert data["model"] == "actor" - assert data["content"] == [{"type": "text", "text": "plain response"}] - assert data["stop_reason"] == "end_turn" - assert data["usage"] == {"input_tokens": 3, "output_tokens": 1} - assert segments[0].prompt_ids == [1, 2, 3] - assert segments[0].response_ids == [581] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_anthropic_messages_endpoint_streams_blocks_and_records_segments(monkeypatch): - async def fake_generate(prompt_ids, session, body, app, **kwargs): - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[601], finish_reason="stop", output_log_probs=[-0.6]) - - async def run_case(): - monkeypatch.setattr(anthropic, "_generate", fake_generate) - raw_output = ( - "delegate inspect" - ) - tokenizer = ToyTokenizer({(601,): raw_output}) - adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") - adapter.open_session("sid-anthropic-stream", sampling_defaults={"max_new_tokens": 8}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - resp = await client.post( - "/v1/messages", - headers={"Authorization": "Bearer sid-anthropic-stream"}, - json={ - "model": "actor", - "system": "be useful", - "stream": True, - "max_tokens": 4, - "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], - "tools": [ - { - "name": "Task", - "description": "spawn subagent", - "input_schema": {"type": "object", "properties": {"description": {"type": "string"}}}, - } - ], - }, - ) - raw = await resp.text() - finally: - await client.close() - - events = _parse_sse(raw) - names = [name for name, _ in events] - starts = [payload for name, payload in events if name == "content_block_start"] - deltas = [payload for name, payload in events if name == "content_block_delta"] - message_delta = next(payload for name, payload in events if name == "message_delta") - segments = await adapter.finish_session("sid-anthropic-stream") - assert resp.status == 200 - assert names[0] == "message_start" - assert names[-1] == "message_stop" - assert any(s["content_block"]["type"] == "text" for s in starts) - assert any(s["content_block"]["type"] == "tool_use" and s["content_block"]["name"] == "Task" for s in starts) - assert any(d["delta"].get("text") == "delegate" for d in deltas) - assert any(json.loads(d["delta"].get("partial_json", "{}")) == {"description": "inspect"} for d in deltas) - assert message_delta["delta"]["stop_reason"] == "tool_use" - assert message_delta["usage"] == {"input_tokens": 3, "output_tokens": 1} - assert segments[0].metadata["segment_kind"] == "final" - assert segments[0].prompt_ids == [1, 2, 3] - assert segments[0].response_ids == [601] - assert segments[0].loss_mask == [1] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_responses_multiturn_uses_vllm_tokens_for_training_segment(): - async def run_case(): - upstream = FakeVLLM( - [ - [(-0.20, 20), (-0.21, 21)], - [(-0.40, 40)], - ] - ) - upstream_app = web.Application() - upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) - upstream_server = TestServer(upstream_app) - await upstream_server.start_server() - - tool_raw = "vime" - tokenizer = ScriptedTokenizer( - prompts=[ - [10, 11], - [10, 11, 20, 21, 30, 31], - ], - outputs={ - (20, 21): tool_raw, - (40,): "done", - }, - ) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url=str(upstream_server.make_url("")).rstrip("/")) - adapter.open_session("sid-openai-token", sampling_defaults={"max_new_tokens": 99}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - first = await client.post( - "/v1/responses", - headers={"Authorization": "Bearer sid-openai-token"}, - json={ - "model": "actor", - "input": "find vime", - "max_output_tokens": 5, - "tools": [ - { - "type": "function", - "name": "lookup", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - first_data = await first.json() - function_call = next(item for item in first_data["output"] if item["type"] == "function_call") - - second = await client.post( - "/v1/responses", - headers={"Authorization": "Bearer sid-openai-token"}, - json={ - "model": "actor", - "input": [ - {"role": "user", "content": "find vime"}, - function_call, - { - "type": "function_call_output", - "call_id": function_call["call_id"], - "output": "found vime", - }, - ], - "max_output_tokens": 7, - "tools": [ - { - "type": "function", - "name": "lookup", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - second_data = await second.json() - finally: - await client.close() - await upstream_server.close() - - segments = await adapter.finish_session("sid-openai-token") - assert first.status == 200 - assert second.status == 200 - assert function_call["name"] == "lookup" - assert function_call["arguments"] == '{"query": "vime"}' - assert second_data["output"][0]["content"][0]["text"] == "done" - assert [req["token_ids"] for req in upstream.requests] == [[10, 11], [10, 11, 20, 21, 30, 31]] - assert upstream.routing_keys == ["sid-openai-token", "sid-openai-token"] - assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 - assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 - assert segments[0].prompt_ids == [10, 11] - assert segments[0].response_ids == [20, 21, 30, 31, 40] - assert segments[0].loss_mask == [1, 1, 0, 0, 1] - assert segments[0].rollout_log_probs == [-0.20, -0.21, 0.0, 0.0, -0.40] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_anthropic_messages_multiturn_uses_vllm_tokens_for_training_segment(): - async def run_case(): - upstream = FakeVLLM( - [ - [(-1.20, 120), (-1.21, 121)], - [(-1.40, 140), (-1.41, 141)], - ] - ) - upstream_app = web.Application() - upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) - upstream_server = TestServer(upstream_app) - await upstream_server.start_server() - - tool_raw = "vime" - tokenizer = ScriptedTokenizer( - prompts=[ - [110, 111], - [110, 111, 120, 121, 130], - ], - outputs={ - (120, 121): tool_raw, - (140, 141): "anthropic done", - }, - ) - adapter = anthropic.AnthropicAdapter( - tokenizer=tokenizer, - vllm_url=str(upstream_server.make_url("")).rstrip("/"), - ) - adapter.open_session("sid-anthropic-token", sampling_defaults={"max_new_tokens": 99}) - client = TestClient(TestServer(adapter.app)) - await client.start_server() - try: - first = await client.post( - "/v1/messages", - headers={"Authorization": "Bearer sid-anthropic-token"}, - json={ - "model": "actor", - "max_tokens": 5, - "messages": [{"role": "user", "content": [{"type": "text", "text": "find vime"}]}], - "tools": [ - { - "name": "lookup", - "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - first_data = await first.json() - tool_use = next(block for block in first_data["content"] if block["type"] == "tool_use") - - second = await client.post( - "/v1/messages", - headers={"Authorization": "Bearer sid-anthropic-token"}, - json={ - "model": "actor", - "max_tokens": 7, - "messages": [ - {"role": "user", "content": [{"type": "text", "text": "find vime"}]}, - {"role": "assistant", "content": first_data["content"]}, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": tool_use["id"], - "content": "found vime", - } - ], - }, - ], - "tools": [ - { - "name": "lookup", - "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ], - }, - ) - second_data = await second.json() - finally: - await client.close() - await upstream_server.close() - - segments = await adapter.finish_session("sid-anthropic-token") - assert first.status == 200 - assert second.status == 200 - assert tool_use["name"] == "lookup" - assert tool_use["input"] == {"query": "vime"} - assert second_data["content"] == [{"type": "text", "text": "anthropic done"}] - assert [req["token_ids"] for req in upstream.requests] == [[110, 111], [110, 111, 120, 121, 130]] - assert upstream.routing_keys == ["sid-anthropic-token", "sid-anthropic-token"] - assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 - assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 - assert segments[0].prompt_ids == [110, 111] - assert segments[0].response_ids == [120, 121, 130, 140, 141] - assert segments[0].loss_mask == [1, 1, 0, 1, 1] - assert segments[0].rollout_log_probs == [-1.20, -1.21, 0.0, -1.40, -1.41] - - asyncio.run(run_case()) - - -@pytest.mark.unit -def test_openai_generate_posts_token_ids_and_extracts_logprobs(): - async def run_case(): - captured = {} - captured_headers = {} - - async def handle_generate(request): - captured_headers.update(request.headers) - captured.update(await request.json()) - return web.json_response( - { - "choices": [ - { - "token_ids": [701, 702], - "logprobs": {"content": [{"logprob": -0.7}, {"logprob": -0.8}]}, - "finish_reason": "stop", - } - ] - } - ) - - upstream_app = web.Application() - upstream_app.router.add_post("/inference/v1/generate", handle_generate) - server = TestServer(upstream_app) - await server.start_server() - try: - session = openai.Session(sampling_defaults={"max_new_tokens": 9}) - turn = await openai._generate( - [11, 12], - session, - {"max_tokens": 3, "temperature": 0.25, "stop": [""]}, - {VLLM_URL_KEY: str(server.make_url("")).rstrip("/")}, - ) - finally: - await server.close() - - assert captured["token_ids"] == [11, 12] - assert captured["sampling_params"]["logprobs"] == 1 - assert captured_headers.get("x-session-id") is None - assert captured["sampling_params"]["max_tokens"] == 3 - assert captured["sampling_params"]["temperature"] == 0.25 - assert captured["sampling_params"]["stop"] == [""] - assert turn.prompt_ids == [11, 12] - assert turn.output_ids == [701, 702] - assert turn.output_log_probs == [-0.7, -0.8] - - asyncio.run(run_case()) - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_sdk_adapters.py b/tests/test_agent_sdk_adapters.py deleted file mode 100644 index d947cd67a..000000000 --- a/tests/test_agent_sdk_adapters.py +++ /dev/null @@ -1,418 +0,0 @@ -import asyncio -import sys -from pathlib import Path - -import httpx -import pytest -from aiohttp.test_utils import TestClient, TestServer - -REPO_ROOT = Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from vime.agent.adapters import anthropic, openai -from vime.agent.trajectory import TurnRecord - - -NUM_GPUS = 0 - - -agents = pytest.importorskip("agents") -anthropic_sdk = pytest.importorskip("anthropic") -openai_sdk = pytest.importorskip("openai") - - -class SDKTokenizer: - def __init__(self, outputs: list[str]) -> None: - self.outputs = outputs - self.rendered: list[tuple[list[dict], list[dict] | None]] = [] - - def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): - self.rendered.append((list(messages), tools)) - return list(range(1, len(messages) + 2)) - - def decode(self, ids, skip_special_tokens=False): - return self.outputs[ids[0] - 1] - - -@pytest.mark.integration -def test_openai_agents_sdk_responses_runs_tool_loop_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), - output_ids=[len(calls)], - finish_reason="stop", - output_log_probs=[-0.1 * len(calls)], - ) - - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = SDKTokenizer( - [ - "vime", - "final after tool", - ] - ) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/v1/")) - http_client = httpx.AsyncClient(trust_env=False) - oai = openai_sdk.AsyncOpenAI( - api_key="sdk-openai", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - @agents.function_tool - def lookup(query: str) -> str: - return f"found {query}" - - agents.set_tracing_disabled(True) - model = agents.OpenAIResponsesModel(model="actor", openai_client=oai) - agent = agents.Agent( - name="sdk-responses", - instructions="Use lookup.", - model=model, - tools=[lookup], - model_settings=agents.ModelSettings(max_tokens=4), - ) - try: - result = await agents.Runner.run(agent, "find vime") - finally: - await client.close() - await oai.close() - - segments = await adapter.finish_session("sdk-openai") - assert result.final_output == "final after tool" - assert len(calls) == 2 - assert calls[0]["body"]["max_output_tokens"] == 4 - assert calls[0]["body"]["tools"][0]["name"] == "lookup" - assert calls[1]["body"]["input"][-1] == { - "call_id": calls[1]["body"]["input"][-2]["call_id"], - "output": "found vime", - "type": "function_call_output", - } - assert tokenizer.rendered[0][0] == [ - {"role": "system", "content": "Use lookup."}, - {"role": "user", "content": "find vime"}, - ] - assert tokenizer.rendered[1][0][-1] == { - "role": "tool", - "content": "found vime", - "tool_call_id": calls[1]["body"]["input"][-2]["call_id"], - } - assert segments[0].metadata["segment_kind"] == "final" - assert segments[0].response_ids[-1] == 2 - assert segments[0].loss_mask[-1] == 1 - - asyncio.run(run_case()) - - -@pytest.mark.integration -def test_openai_agents_sdk_chat_completions_runs_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.2] - ) - - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = SDKTokenizer(["chat final"]) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/v1/")) - http_client = httpx.AsyncClient(trust_env=False) - oai = openai_sdk.AsyncOpenAI( - api_key="sdk-openai-chat", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - agents.set_tracing_disabled(True) - model = agents.OpenAIChatCompletionsModel(model="actor", openai_client=oai) - agent = agents.Agent( - name="sdk-chat", - instructions="Be short.", - model=model, - model_settings=agents.ModelSettings(max_tokens=5), - ) - try: - result = await agents.Runner.run(agent, "say hi") - finally: - await client.close() - await oai.close() - - segments = await adapter.finish_session("sdk-openai-chat") - assert result.final_output == "chat final" - assert calls[0]["body"]["max_tokens"] == 5 - assert calls[0]["body"]["messages"] == [ - {"content": "Be short.", "role": "system"}, - {"role": "user", "content": "say hi"}, - ] - assert segments[0].prompt_ids == [1, 2, 3] - assert segments[0].response_ids == [1] - assert segments[0].loss_mask == [1] - - asyncio.run(run_case()) - - -@pytest.mark.integration -def test_openai_sdk_chat_completion_streaming_runs_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.25] - ) - - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = SDKTokenizer( - ["streamed via sdk vime"] - ) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/v1/")) - http_client = httpx.AsyncClient(trust_env=False) - oai = openai_sdk.AsyncOpenAI( - api_key="sdk-openai-chat-stream", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - try: - stream = await oai.chat.completions.create( - model="actor", - messages=[{"role": "user", "content": "call lookup"}], - tools=[ - { - "type": "function", - "function": { - "name": "lookup", - "description": "search", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, - }, - } - ], - stream=True, - ) - text_parts = [] - tool_names = [] - tool_arguments = [] - finish_reasons = [] - usages = [] - async for chunk in stream: - choice = chunk.choices[0] - if choice.delta.content: - text_parts.append(choice.delta.content) - if choice.delta.tool_calls: - for tool_call in choice.delta.tool_calls: - tool_names.append(tool_call.function.name) - tool_arguments.append(tool_call.function.arguments) - if choice.finish_reason: - finish_reasons.append(choice.finish_reason) - if chunk.usage: - usages.append(chunk.usage) - finally: - await client.close() - await oai.close() - - segments = await adapter.finish_session("sdk-openai-chat-stream") - assert "".join(text_parts) == "streamed via sdk" - assert tool_names == ["lookup"] - assert tool_arguments == ['{"query": "vime"}'] - assert finish_reasons == ["tool_calls"] - assert usages[-1].prompt_tokens == 2 - assert usages[-1].completion_tokens == 1 - assert calls[0]["body"]["stream"] is True - assert segments[0].response_ids == [1] - - asyncio.run(run_case()) - - -@pytest.mark.integration -def test_openai_sdk_responses_streaming_runs_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.35] - ) - - monkeypatch.setattr(openai, "_generate", fake_generate) - tokenizer = SDKTokenizer(["response stream via sdk"]) - adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/v1/")) - http_client = httpx.AsyncClient(trust_env=False) - oai = openai_sdk.AsyncOpenAI( - api_key="sdk-openai-responses-stream", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - try: - stream = await oai.responses.create( - model="actor", - instructions="Be brief.", - input="say hi", - stream=True, - ) - event_types = [] - deltas = [] - completed_response = None - async for event in stream: - event_types.append(event.type) - if event.type == "response.output_text.delta": - deltas.append(event.delta) - if event.type == "response.completed": - completed_response = event.response - finally: - await client.close() - await oai.close() - - segments = await adapter.finish_session("sdk-openai-responses-stream") - assert event_types == ["response.created", "response.output_text.delta", "response.completed"] - assert "".join(deltas) == "response stream via sdk" - assert completed_response.status == "completed" - assert completed_response.usage.input_tokens == 3 - assert completed_response.usage.output_tokens == 1 - assert calls[0]["body"]["stream"] is True - assert segments[0].response_ids == [1] - - asyncio.run(run_case()) - - -@pytest.mark.integration -def test_anthropic_sdk_non_streaming_messages_runs_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.28] - ) - - monkeypatch.setattr(anthropic, "_generate", fake_generate) - tokenizer = SDKTokenizer(["anthropic json"]) - adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/")) - http_client = httpx.AsyncClient(trust_env=False) - anth = anthropic_sdk.AsyncAnthropic( - api_key="sdk-anthropic-json", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - try: - message = await anth.messages.create( - model="actor", - max_tokens=6, - system="Be direct.", - messages=[{"role": "user", "content": "say hi"}], - ) - finally: - await client.close() - await anth.close() - - segments = await adapter.finish_session("sdk-anthropic-json") - assert message.type == "message" - assert message.content[0].type == "text" - assert message.content[0].text == "anthropic json" - assert message.stop_reason == "end_turn" - assert message.usage.input_tokens == 3 - assert message.usage.output_tokens == 1 - assert calls[0]["body"]["max_tokens"] == 6 - assert segments[0].response_ids == [1] - - asyncio.run(run_case()) - - -@pytest.mark.integration -def test_anthropic_sdk_streaming_messages_runs_against_adapter(monkeypatch): - async def run_case(): - calls = [] - - async def fake_generate(prompt_ids, session, body, app, **kwargs): - calls.append({"prompt_ids": list(prompt_ids), "body": body}) - return TurnRecord( - prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.3] - ) - - monkeypatch.setattr(anthropic, "_generate", fake_generate) - tokenizer = SDKTokenizer(["anthropic final"]) - adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") - client = TestClient(TestServer(adapter.app)) - await client.start_server() - base_url = str(client.make_url("/")) - http_client = httpx.AsyncClient(trust_env=False) - anth = anthropic_sdk.AsyncAnthropic( - api_key="sdk-anthropic", - base_url=base_url, - max_retries=0, - http_client=http_client, - ) - - try: - stream = await anth.messages.create( - model="actor", - max_tokens=6, - system="Be direct.", - messages=[{"role": "user", "content": "say hi"}], - stream=True, - ) - text_parts = [] - event_types = [] - async for event in stream: - event_types.append(event.type) - if event.type == "content_block_delta" and getattr(event.delta, "text", None): - text_parts.append(event.delta.text) - finally: - await client.close() - await anth.close() - - segments = await adapter.finish_session("sdk-anthropic") - assert "".join(text_parts) == "anthropic final" - assert event_types == [ - "message_start", - "content_block_start", - "content_block_delta", - "content_block_stop", - "message_delta", - "message_stop", - ] - assert calls[0]["body"]["max_tokens"] == 6 - assert tokenizer.rendered[0][0] == [ - {"role": "system", "content": "Be direct."}, - {"role": "user", "content": "say hi"}, - ] - assert segments[0].prompt_ids == [1, 2, 3] - assert segments[0].response_ids == [1] - assert segments[0].loss_mask == [1] - - asyncio.run(run_case()) - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_trajectory.py b/tests/test_agent_trajectory.py deleted file mode 100644 index 3abf88610..000000000 --- a/tests/test_agent_trajectory.py +++ /dev/null @@ -1,143 +0,0 @@ -import sys -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from vime.agent.trajectory import TurnRecord, TurnSegment, merge_turn_segments, merge_turns - - -NUM_GPUS = 0 - - -def _turn(prompt_ids: list[int], output_ids: list[int], output_log_probs: list[float] | None = None) -> TurnRecord: - return TurnRecord( - prompt_ids=prompt_ids, - output_ids=output_ids, - finish_reason="stop", - output_log_probs=( - output_log_probs if output_log_probs is not None else [-token_id / 100 for token_id in output_ids] - ), - ) - - -@pytest.mark.unit -def test_merge_turns_preserves_matched_prefix_on_prompt_drift(): - segment = merge_turns( - [ - _turn([10], [11]), - _turn([10, 11, 21], [12]), - _turn([10, 11, 21, 12, 31], [13]), - _turn([10, 11, 21, 12, 22], [14]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [10] - assert segment.response_ids == [11, 21, 12, 22, 14] - assert segment.loss_mask == [1, 0, 1, 0, 1] - assert segment.rollout_log_probs == [-0.11, 0.0, -0.12, 0.0, -0.14] - - -@pytest.mark.unit -def test_merge_turns_drops_middle_turn_when_next_prompt_skips_it(): - segment = merge_turns( - [ - _turn([10], [11]), - _turn([10, 11, 21], [12]), - _turn([10, 11, 22], [13]), - _turn([10, 11, 22, 13, 31], [14]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [10] - assert segment.response_ids == [11, 22, 13, 31, 14] - assert segment.loss_mask == [1, 0, 1, 0, 1] - assert segment.rollout_log_probs == [-0.11, 0.0, -0.13, 0.0, -0.14] - - -@pytest.mark.unit -def test_merge_turns_handles_consecutive_prompt_drifts(): - segment = merge_turns( - [ - _turn([10], [11]), - _turn([10, 11, 21], [12]), - _turn([10, 11, 22], [13]), - _turn([10, 11, 23], [14]), - _turn([10, 11, 23, 14, 31], [15]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [10] - assert segment.response_ids == [11, 23, 14, 31, 15] - assert segment.loss_mask == [1, 0, 1, 0, 1] - assert segment.rollout_log_probs == [-0.11, 0.0, -0.14, 0.0, -0.15] - - -@pytest.mark.unit -def test_merge_turns_masks_whole_output_when_prompt_drift_splits_it(): - segment = merge_turns( - [ - _turn([10], [11, 12, 13, 14]), - _turn([10, 11, 12, 99, 14], [15]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [10] - assert segment.response_ids == [11, 12, 99, 14, 15] - assert segment.loss_mask == [0, 0, 0, 0, 1] - assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, -0.15] - - -@pytest.mark.unit -def test_merge_turns_masks_whole_output_when_prompt_drift_changes_token_count(): - segment = merge_turns( - [ - _turn([10], [11, 12, 13, 14]), - _turn([10, 11, 12, 99, 100, 14], [15]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [10] - assert segment.response_ids == [11, 12, 99, 100, 14, 15] - assert segment.loss_mask == [0, 0, 0, 0, 0, 1] - assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, 0.0, -0.15] - - -@pytest.mark.unit -def test_merge_turns_restarts_when_prompt_base_changes(): - segment = merge_turns( - [ - _turn([10], [11]), - _turn([20, 21], [22]), - _turn([20, 21, 22, 23], [24]), - ] - ) - - assert segment is not None - assert segment.prompt_ids == [20, 21] - assert segment.response_ids == [22, 23, 24] - assert segment.loss_mask == [1, 0, 1] - assert segment.rollout_log_probs == [-0.22, 0.0, -0.24] - - -@pytest.mark.unit -def test_merge_turn_segments_keeps_oversized_segments(): - segments = [TurnSegment(turns=[_turn([10, 11, 12], [13, 14])])] - - merged = merge_turn_segments(segments) - - assert len(merged) == 1 - assert merged[0].prompt_ids == [10, 11, 12] - assert merged[0].response_ids == [13, 14] - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_block_fp8_zero_block.py b/tests/test_block_fp8_zero_block.py new file mode 100644 index 000000000..16885d25c --- /dev/null +++ b/tests/test_block_fp8_zero_block.py @@ -0,0 +1,79 @@ +"""CPU unit tests for block-wise FP8 quantization in tools/convert_hf_to_fp8.py. + +Pins the contract that an all-zero quantization block must not produce +NaN weights. ``block_fp8`` (the default quantization strategy) computed +``scale = block_max / FP8_MAX`` without clamping the block max away from +zero, unlike ``channel_fp8`` and ``tensor_fp8`` which both clamp to +1e-12. An all-zero block (padding rows, an unused MoE expert, ...) gave +``scale == 0`` and ``qweight == 0 / 0 == NaN``, silently writing NaN +weights into the converted checkpoint. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +pytest.importorskip("safetensors") +torch = pytest.importorskip("torch") + +NUM_GPUS = 0 + + +def _load_converter(): + module_path = Path(__file__).resolve().parents[1] / "tools" / "convert_hf_to_fp8.py" + spec = importlib.util.spec_from_file_location("convert_hf_to_fp8", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def converter(): + return _load_converter() + + +@pytest.mark.unit +def test_block_fp8_all_zero_block_has_no_nan(converter): + # First 128x128 tile is non-zero, the other three are all-zero blocks. + weight = torch.zeros(256, 256, dtype=torch.bfloat16) + weight[0, 0] = 1.0 + + qweight, scale = converter.block_fp8(weight, (128, 128)) + + assert not torch.isnan(qweight.float()).any() + assert not torch.isinf(qweight.float()).any() + # No scale may be zero: dequantization multiplies by the scale, and a + # zero scale is exactly what turned the all-zero block into NaN. + assert (scale > 0).all() + + +@pytest.mark.unit +def test_block_fp8_zero_block_roundtrips_to_zero(converter): + weight = torch.zeros(128, 128, dtype=torch.bfloat16) + + qweight, scale = converter.block_fp8(weight, (128, 128)) + + dequantized = qweight.float() * scale.float() + assert (dequantized == 0).all() + + +@pytest.mark.unit +def test_block_fp8_nonzero_blocks_unaffected(converter): + torch.manual_seed(0) + weight = torch.randn(256, 256, dtype=torch.float32).to(torch.bfloat16) + + qweight, scale = converter.block_fp8(weight, (128, 128)) + + assert not torch.isnan(qweight.float()).any() + dequantized = qweight.float() * scale.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1).float() + max_err = (dequantized - weight.float()).abs().max() + # FP8 e4m3 relative error is ~2^-3, so a tolerance of 0.5 is generous + # for randn-scale values and only guards against gross corruption. + assert max_err < 0.5 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_chunked_gae.py b/tests/test_chunked_gae.py deleted file mode 100644 index a57ad6217..000000000 --- a/tests/test_chunked_gae.py +++ /dev/null @@ -1,63 +0,0 @@ -import time -import pytest -import torch - -from vime.utils.ppo_utils import chunked_gae, vanilla_gae - - -@pytest.mark.parametrize( - "B,T", - [ - (16, 4096), - (32, 8192), - (256, 128 * 1024), - ], -) -@pytest.mark.parametrize("chunk_size", [64, 128, 256]) -def test_gae_parallel_matches_serial(B, T, chunk_size): - """ - Test that chunked_gae (parallel-scan) matches vanilla_gae (batch-serial) - under various shapes, chunk sizes and dtypes. - """ - device = "cuda" if torch.cuda.is_available() else "cpu" - torch.manual_seed(0) - - rewards = torch.randn(B, T, device=device, dtype=torch.float32) - values = torch.randn(B, T, device=device, dtype=torch.float32) - - gamma, lam = 0.99, 0.95 - - # ---------- Serial ---------- - if device == "cuda": - torch.cuda.synchronize() - t0 = time.time() - adv_s, ret_s = vanilla_gae(rewards, values, gamma, lam) - if device == "cuda": - torch.cuda.synchronize() - t1 = time.time() - serial_time = t1 - t0 - - # ---------- Parallel-scan ---------- - if device == "cuda": - torch.cuda.synchronize() - t0 = time.time() - adv_p, ret_p = chunked_gae(rewards, values, gamma, lam, chunk_size=chunk_size) - if device == "cuda": - torch.cuda.synchronize() - t1 = time.time() - parallel_time = t1 - t0 - - # ---------- Accuracy ---------- - adv_err = (adv_s - adv_p).abs().max().item() - ret_err = (ret_s - ret_p).abs().max().item() - - atol = 1e-5 - assert adv_err < atol, f"adv error too large: {adv_err}" - assert ret_err < atol, f"ret error too large: {ret_err}" - - # ---------- logging ---------- - print(f"\n[GAE Test] B={B}, T={T}, chunk={chunk_size}") - print(f" Serial : {serial_time:.6f} s") - print(f" Parallel : {parallel_time:.6f} s") - print(f" Speedup : x{serial_time / parallel_time:.2f}") - print(f" Max diff adv={adv_err:.3e}, ret={ret_err:.3e}") diff --git a/tests/test_cispo_loss.py b/tests/test_cispo_loss.py new file mode 100644 index 000000000..6333ce3f6 --- /dev/null +++ b/tests/test_cispo_loss.py @@ -0,0 +1,52 @@ +"""CPU tests for compute_cispo_loss (MiniMax-M1, https://arxiv.org/abs/2506.13585).""" + +import math + +import pytest +import torch + +from vime.utils.ppo_utils import compute_cispo_loss + +NUM_GPUS = 0 + +ADVANTAGES = torch.tensor([1.0, -0.5, 2.0, -1.0]) +LOG_PROBS = torch.tensor([-0.7, -1.2, -0.4, -2.1]) + +# (eps_clip, eps_clip_high, raw IS ratios, ratios after clamp to [1 - eps_clip, 1 + eps_clip_high]) +CLIP_CASES = [ + pytest.param(0.2, 0.28, [1.0, 1.14, 1.56, 0.4], [1.0, 1.14, 1.28, 0.8], id="ppo_band"), + pytest.param(1.0, 4.0, [1.0, 3.0, 9.0, 0.4], [1.0, 3.0, 5.0, 0.4], id="wide_minimax_band"), +] + + +@pytest.mark.parametrize("eps_clip, eps_clip_high, ratios, clamped", CLIP_CASES) +def test_compute_cispo_loss_matches_closed_form_surrogate(eps_clip, eps_clip_high, ratios, clamped): + ppo_kl = -torch.tensor([math.log(r) for r in ratios]) + + pg_losses, clipfrac = compute_cispo_loss(ppo_kl, LOG_PROBS, ADVANTAGES, eps_clip, eps_clip_high) + + expected_losses = -torch.tensor(clamped) * ADVANTAGES * LOG_PROBS + torch.testing.assert_close(pg_losses, expected_losses, rtol=1e-6, atol=1e-6) + expected_clipfrac = torch.tensor([float(c != r) for c, r in zip(clamped, ratios, strict=True)]) + torch.testing.assert_close(clipfrac, expected_clipfrac) + + +@pytest.mark.parametrize("eps_clip, eps_clip_high, ratios, clamped", CLIP_CASES) +def test_compute_cispo_loss_gradient_flows_only_through_log_probs(eps_clip, eps_clip_high, ratios, clamped): + # ratio = exp(-ppo_kl) = exp(log_ratios): if CISPO failed to stop-gradient the + # clipped IS ratio, backward would populate log_ratios.grad. + log_ratios = torch.tensor([math.log(r) for r in ratios], requires_grad=True) + ppo_kl = -log_ratios + log_probs = LOG_PROBS.clone().requires_grad_() + + pg_losses, _ = compute_cispo_loss(ppo_kl, log_probs, ADVANTAGES, eps_clip, eps_clip_high) + pg_losses.sum().backward() + + torch.testing.assert_close(log_probs.grad, -torch.tensor(clamped) * ADVANTAGES, rtol=1e-6, atol=1e-6) + assert log_ratios.grad is None or torch.all( + log_ratios.grad == 0 + ), f"CISPO must stop-gradient on the IS ratio; log_ratios.grad={log_ratios.grad}" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_deep_ep_tms_patch.py b/tests/test_deep_ep_tms_patch.py new file mode 100644 index 000000000..01fad32e2 --- /dev/null +++ b/tests/test_deep_ep_tms_patch.py @@ -0,0 +1,104 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + +NUM_GPUS = 0 + + +def _load_megatron_utils_init(monkeypatch, buffer_cls, tms_impl, module_name): + deep_ep = types.ModuleType("deep_ep") + deep_ep.Buffer = buffer_cls + monkeypatch.setitem(sys.modules, "deep_ep", deep_ep) + + torch_memory_saver_module = types.ModuleType("torch_memory_saver") + torch_memory_saver_module.torch_memory_saver = types.SimpleNamespace(_impl=tms_impl) + monkeypatch.setitem(sys.modules, "torch_memory_saver", torch_memory_saver_module) + monkeypatch.setitem(sys.modules, "megatron", types.ModuleType("megatron")) + + package_path = Path(__file__).parents[1] / "vime" / "backends" / "megatron_utils" + spec = importlib.util.spec_from_file_location( + module_name, + package_path / "__init__.py", + submodule_search_locations=[str(package_path)], + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + + +@pytest.mark.unit +def test_deep_ep_init_restores_original_tms_region(monkeypatch): + events = [] + + class FakeCdll: + def __init__(self): + self.interesting_region = False + + def tms_get_interesting_region(self): + events.append(("get", self.interesting_region)) + return self.interesting_region + + def tms_set_interesting_region(self, enabled): + self.interesting_region = enabled + events.append(("set", enabled)) + + cdll = FakeCdll() + + class FakeBuffer: + def __init__(self): + events.append(("init", cdll.interesting_region)) + + tms_impl = types.SimpleNamespace(_binary_wrapper=types.SimpleNamespace(cdll=cdll)) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: events.append(("sync", cdll.interesting_region))) + + _load_megatron_utils_init(monkeypatch, FakeBuffer, tms_impl, "_test_megatron_utils_tms_restore") + FakeBuffer() + + assert events == [ + ("get", False), + ("set", False), + ("init", False), + ("sync", False), + ("set", False), + ] + assert cdll.interesting_region is False + + +@pytest.mark.unit +def test_deep_ep_init_restores_tms_region_after_failure(monkeypatch): + events = [] + + class FakeCdll: + interesting_region = True + + def tms_get_interesting_region(self): + return self.interesting_region + + def tms_set_interesting_region(self, enabled): + self.interesting_region = enabled + events.append(("set", enabled)) + + cdll = FakeCdll() + + class FailingBuffer: + def __init__(self): + events.append(("init", cdll.interesting_region)) + raise RuntimeError("buffer init failed") + + tms_impl = types.SimpleNamespace(_binary_wrapper=types.SimpleNamespace(cdll=cdll)) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: events.append(("sync", cdll.interesting_region))) + + _load_megatron_utils_init(monkeypatch, FailingBuffer, tms_impl, "_test_megatron_utils_tms_failure") + with pytest.raises(RuntimeError, match="buffer init failed"): + FailingBuffer() + + assert events == [("set", False), ("init", False), ("set", True)] + assert cdll.interesting_region is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_discounted_returns.py b/tests/test_discounted_returns.py new file mode 100644 index 000000000..f776a3de5 --- /dev/null +++ b/tests/test_discounted_returns.py @@ -0,0 +1,108 @@ +import sys +import types + +import pytest +import torch + +from vime.utils.ppo_utils import chunked_discounted_returns, chunked_gae, get_reinforce_plus_plus_returns, vanilla_gae + + +NUM_GPUS = 0 + + +def _serial_discounted_returns(rewards: torch.Tensor, discount: float) -> torch.Tensor: + returns = torch.zeros_like(rewards) + running_return = torch.zeros(rewards.size(0), device=rewards.device, dtype=rewards.dtype) + for t in reversed(range(rewards.size(1))): + running_return = rewards[:, t] + discount * running_return + returns[:, t] = running_return + return returns + + +@pytest.mark.parametrize("discount", [0.0, 0.5, 0.99, 1.0]) +@pytest.mark.parametrize("batch_size,sequence_length", [(1, 1), (3, 127), (3, 128), (3, 129), (3, 1000)]) +@pytest.mark.parametrize( + "dtype,atol,rtol", + [ + (torch.float32, 1e-4, 1e-5), + (torch.float64, 1e-10, 1e-10), + ], +) +def test_chunked_discounted_returns_matches_serial(discount, batch_size, sequence_length, dtype, atol, rtol): + torch.manual_seed(0) + rewards = torch.randn(batch_size, sequence_length, dtype=dtype) + + expected = _serial_discounted_returns(rewards, discount) + actual = chunked_discounted_returns(rewards, discount) + + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + assert actual.dtype == rewards.dtype + assert actual.device == rewards.device + + +def test_chunked_discounted_returns_preserves_right_padding(): + torch.manual_seed(0) + lengths = [3, 129, 511] + rewards = torch.zeros(len(lengths), max(lengths)) + for i, length in enumerate(lengths): + rewards[i, :length] = torch.randn(length) + + actual = chunked_discounted_returns(rewards, 0.99) + + for i, length in enumerate(lengths): + expected = _serial_discounted_returns(rewards[i : i + 1, :length], 0.99)[0] + torch.testing.assert_close(actual[i, :length], expected, atol=1e-4, rtol=1e-5) + assert torch.count_nonzero(actual[i, length:]) == 0 + + +def test_reinforce_plus_plus_returns_matches_serial_for_variable_lengths(monkeypatch): + mpu = types.SimpleNamespace(get_context_parallel_world_size=lambda: 1) + megatron = types.ModuleType("megatron") + megatron_core = types.ModuleType("megatron.core") + megatron_core.mpu = mpu + megatron.core = megatron_core + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", megatron_core) + + rewards = torch.tensor([2.0, -1.0]) + kl = [torch.tensor([0.2, 0.1, 0.3, 0.4, 0.5]), torch.tensor([0.4, 0.2, 0.1])] + loss_masks = [torch.tensor([0.0, 1.0, 1.0, 1.0, 0.0]), torch.tensor([1.0, 1.0, 1.0])] + kl_coef = 0.1 + gamma = 0.99 + + actual = get_reinforce_plus_plus_returns( + rewards=rewards, + kl=kl, + loss_masks=loss_masks, + response_lengths=[3, 3], + total_lengths=[5, 3], + kl_coef=kl_coef, + gamma=gamma, + ) + + expected = [] + for reward, kl_for_seq, mask in zip(rewards, kl, loss_masks, strict=True): + token_rewards = -kl_coef * kl_for_seq * mask + token_rewards[mask.nonzero(as_tuple=True)[0][-1]] += reward + expected.append(_serial_discounted_returns(token_rewards.unsqueeze(0), gamma)[0]) + + assert len(actual) == len(expected) + for actual_for_seq, expected_for_seq in zip(actual, expected, strict=True): + torch.testing.assert_close(actual_for_seq, expected_for_seq, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("sequence_length", [127, 128, 129]) +def test_chunked_gae_matches_serial_after_scan_reuse(sequence_length): + torch.manual_seed(0) + rewards = torch.randn(3, sequence_length) + values = torch.randn(3, sequence_length) + + expected = vanilla_gae(rewards, values, gamma=0.99, lambd=0.95) + actual = chunked_gae(rewards, values, gamma=0.99, lambd=0.95) + + torch.testing.assert_close(actual[0], expected[0], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual[1], expected[1], atol=1e-5, rtol=1e-5) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_docs_consistency.py b/tests/test_docs_consistency.py new file mode 100644 index 000000000..217a29884 --- /dev/null +++ b/tests/test_docs_consistency.py @@ -0,0 +1,90 @@ +"""Guard documentation references that can be checked without a GPU runtime.""" + +from __future__ import annotations + +import re +from pathlib import Path +from urllib.parse import unquote + +import pytest + +NUM_GPUS = 0 + +ROOT = Path(__file__).resolve().parents[1] +DOC_SOURCES = ( + ROOT / "README.md", + ROOT / "README_zh.md", + ROOT / "docs" / "en", + ROOT / "docs" / "zh", + ROOT / "examples", + ROOT / "docker", +) +MARKDOWN_LINK_RE = re.compile(r"!?\[[^\]]*\]\(([^)\n]+)\)") +COMMAND_PATH_RE = re.compile(r"\b(?:bash|python)\s+((?:scripts?|tests)/[A-Za-z0-9_.+/-]+\.(?:sh|py))") +LOCAL_ANCHOR_RE = re.compile(r"\]\(#([A-Za-z0-9_-]+)\)") +HEADING_RE = re.compile(r"^#{1,6}\s+(.+?)\s*$", re.MULTILINE) + + +def _markdown_files(): + for source in DOC_SOURCES: + if source.is_file(): + yield source + else: + for path in sorted(source.rglob("*.md")): + if "_examples_synced" not in path.parts: + yield path + + +def _local_link_target(raw_target: str) -> str | None: + target = raw_target.strip() + if target.startswith("<") and ">" in target: + target = target[1 : target.index(">")] + else: + target = target.split(maxsplit=1)[0] + + if target.startswith(("#", "/")) or "://" in target or target.startswith(("mailto:", "data:")): + return None + + return unquote(target.split("#", 1)[0].split("?", 1)[0]) or None + + +def test_local_markdown_links_exist(): + missing = [] + for markdown_file in _markdown_files(): + for raw_target in MARKDOWN_LINK_RE.findall(markdown_file.read_text(encoding="utf-8")): + target = _local_link_target(raw_target) + if ( + target is not None + and "_examples_synced" not in Path(target).parts + and not (markdown_file.parent / target).exists() + ): + missing.append(f"{markdown_file.relative_to(ROOT)} -> {target}") + + assert not missing, "Broken local Markdown links:\n" + "\n".join(missing) + + +def test_documented_script_and_test_commands_exist(): + missing = [] + for markdown_file in _markdown_files(): + text = markdown_file.read_text(encoding="utf-8") + for command_path in COMMAND_PATH_RE.findall(text): + if not (ROOT / command_path).is_file(): + missing.append(f"{markdown_file.relative_to(ROOT)} -> {command_path}") + + assert not missing, "Documented command paths do not exist:\n" + "\n".join(missing) + + +@pytest.mark.parametrize("language", ["en", "zh"]) +def test_customization_anchor_links_exist(language): + text = (ROOT / "docs" / language / "get_started" / "customization.md").read_text(encoding="utf-8") + headings = set() + for heading in HEADING_RE.findall(text): + heading = re.sub(r"[^\w\s-]", "", heading.replace("`", "").lower()) + headings.add(re.sub(r"\s+", "-", heading).strip("-")) + + missing = sorted(set(LOCAL_ANCHOR_RE.findall(text)) - headings) + assert not missing, f"Local anchors without matching headings: {missing}" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_dp_schedule.py b/tests/test_dp_schedule.py index edb27add1..0900567c0 100644 --- a/tests/test_dp_schedule.py +++ b/tests/test_dp_schedule.py @@ -20,12 +20,22 @@ def make_args( use_dynamic_batch_size=False, max_tokens_per_gpu=None, balance_data=False, + balance_by_flops=False, ): return SimpleNamespace( micro_batch_size=micro_batch_size, use_dynamic_batch_size=use_dynamic_batch_size, max_tokens_per_gpu=max_tokens_per_gpu, balance_data=balance_data, + balance_by_flops=balance_by_flops, + hidden_size=16, + num_attention_heads=2, + num_query_groups=2, + vocab_size=32, + ffn_hidden_size=64, + num_experts=None, + num_layers=2, + kv_channels=8, ) @@ -131,6 +141,31 @@ def test_static_balance_multi_step(): ) +@pytest.mark.unit +def test_balance_data_distributes_by_flops(): + """balance_data uses FLOPs weights for rank assignment, not raw token sums.""" + total_lengths = [1, 2, 3, 4, 5, 7, 9, 10] + rollout_indices = list(range(8)) + args = make_args(micro_batch_size=1, balance_data=True) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert partitions == [[1, 2, 5, 6], [0, 3, 4, 7]] + assert nmb == [4] + assert gbs_per_step == [8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + ) + + @pytest.mark.unit def test_dynamic_uniform(): """Dynamic mbs on uniform-length samples.""" diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py new file mode 100644 index 000000000..4162d3003 --- /dev/null +++ b/tests/test_empty_colocated_weight_bucket.py @@ -0,0 +1,255 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +NUM_GPUS = 0 + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +class _FakeRemoteMethod: + def __init__(self): + self.calls = [] + + def remote(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return f"ref-{len(self.calls)}" + + +class _FakeEngine: + def __init__(self): + self.update_weights = _FakeRemoteMethod() + + +def _install_fake_deps(monkeypatch): + dist_state = types.SimpleNamespace(rank=0, world_size=2, gathered=None, local_object=None) + + vime_pkg = types.ModuleType("vime") + vime_pkg.__path__ = [str(REPO_ROOT / "vime")] + vime_backends_pkg = types.ModuleType("vime.backends") + vime_backends_pkg.__path__ = [str(REPO_ROOT / "vime" / "backends")] + megatron_utils_pkg = types.ModuleType("vime.backends.megatron_utils") + megatron_utils_pkg.__path__ = [str(REPO_ROOT / "vime" / "backends" / "megatron_utils")] + update_weight_pkg = types.ModuleType("vime.backends.megatron_utils.update_weight") + update_weight_pkg.__path__ = [str(REPO_ROOT / "vime" / "backends" / "megatron_utils" / "update_weight")] + vime_utils_pkg = types.ModuleType("vime.utils") + vime_utils_pkg.__path__ = [str(REPO_ROOT / "vime" / "utils")] + accelerator_mod = types.ModuleType("vime.utils.accelerator") + accelerator_mod.device = lambda: "cuda:0" + accelerator_mod.current_device = lambda: "cuda:0" + accelerator_mod.ipc_collect = lambda: None + + platform_mod = types.ModuleType("vime.platforms") + platform_mod.current_platform = lambda: types.SimpleNamespace(is_npu=False) + + dist_mod = types.ModuleType("torch.distributed") + + def gather_object(obj, object_gather_list, dst, group): + dist_state.local_object = obj + if object_gather_list is not None: + object_gather_list[:] = dist_state.gathered(obj) + + dist_mod.get_rank = lambda: dist_state.rank + dist_mod.get_world_size = lambda group=None: dist_state.world_size + dist_mod.gather_object = gather_object + + torch_mod = types.ModuleType("torch") + torch_mod.Tensor = object + torch_mod.dtype = object + torch_mod.uint8 = "uint8" + torch_mod.distributed = dist_mod + torch_mod.empty = lambda size, dtype, device: {"size": size, "dtype": dtype, "device": device} + torch_mod.no_grad = lambda: (lambda fn: fn) + torch_mod.cuda = types.SimpleNamespace( + current_device=lambda: 0, + get_device_properties=lambda _device: types.SimpleNamespace(uuid="gpu-0"), + ipc_collect=lambda: None, + ) + torch_mod.nn = types.SimpleNamespace(Module=object) + + ray_mod = types.ModuleType("ray") + ray_mod.ObjectRef = object + ray_actor_mod = types.ModuleType("ray.actor") + ray_actor_mod.ActorHandle = object + + mpu_mod = types.ModuleType("megatron.core.mpu") + megatron_mod = types.ModuleType("megatron") + megatron_core_mod = types.ModuleType("megatron.core") + megatron_core_mod.mpu = mpu_mod + + update_weight_common_mod = types.ModuleType("vime.backends.megatron_utils.update_weight.common") + update_weight_common_mod.HfWeightSource = object + update_weight_common_mod.VimeRayWeightSyncClient = object + update_weight_common_mod.create_nccl_trainer = lambda *args, **kwargs: None + + megatron_to_hf_mod = types.ModuleType("vime.backends.megatron_utils.megatron_to_hf") + megatron_to_hf_mod.convert_to_hf = lambda *args, **kwargs: [] + + expert_routing_mod = types.ModuleType("vime.backends.megatron_utils.update_weight.expert_routing") + expert_routing_mod.configure_expert_routing = lambda *args, **kwargs: (None, []) + + hf_weight_iterator_direct_mod = types.ModuleType( + "vime.backends.megatron_utils.update_weight.hf_weight_iterator_direct" + ) + hf_weight_iterator_direct_mod.HfWeightIteratorDirect = lambda *args, **kwargs: None + + vime_utils_types_mod = types.ModuleType("vime.utils.types") + vime_utils_types_mod.ParamInfo = type("ParamInfo", (), {}) + + distributed_utils_mod = types.ModuleType("vime.utils.distributed_utils") + distributed_utils_mod.get_gloo_group = lambda: object() + + update_from_distributed_mod = types.ModuleType( + "vime.backends.megatron_utils.update_weight.update_weight_from_distributed" + ) + update_from_distributed_mod.connect_rollout_engines_from_distributed = lambda *args, **kwargs: None + update_from_distributed_mod.disconnect_rollout_engines_from_distributed = lambda *args, **kwargs: None + update_from_distributed_mod.post_process_weights = lambda *args, **kwargs: None + update_from_distributed_mod.update_weights_from_distributed = lambda *args, **kwargs: [] + + monkeypatch.setitem(sys.modules, "vime", vime_pkg) + monkeypatch.setitem(sys.modules, "vime.platforms", platform_mod) + monkeypatch.setitem(sys.modules, "vime.backends", vime_backends_pkg) + monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils", megatron_utils_pkg) + monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.update_weight", update_weight_pkg) + monkeypatch.setitem(sys.modules, "vime.utils", vime_utils_pkg) + monkeypatch.setitem(sys.modules, "vime.utils.accelerator", accelerator_mod) + monkeypatch.setitem(sys.modules, "torch", torch_mod) + monkeypatch.setitem(sys.modules, "torch.distributed", dist_mod) + monkeypatch.setitem(sys.modules, "ray", ray_mod) + monkeypatch.setitem(sys.modules, "ray.actor", ray_actor_mod) + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.core", megatron_core_mod) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu_mod) + monkeypatch.setitem( + sys.modules, + "vime.backends.megatron_utils.update_weight.common", + update_weight_common_mod, + ) + monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.megatron_to_hf", megatron_to_hf_mod) + monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.update_weight.expert_routing", expert_routing_mod) + monkeypatch.setitem( + sys.modules, + "vime.backends.megatron_utils.update_weight.hf_weight_iterator_direct", + hf_weight_iterator_direct_mod, + ) + monkeypatch.setitem(sys.modules, "vime.utils.types", vime_utils_types_mod) + monkeypatch.setitem(sys.modules, "vime.utils.distributed_utils", distributed_utils_mod) + monkeypatch.setitem( + sys.modules, + "vime.backends.megatron_utils.update_weight.update_weight_from_distributed", + update_from_distributed_mod, + ) + + return dist_state + + +def _load_update_weight_module(monkeypatch): + dist_state = _install_fake_deps(monkeypatch) + + module_name = "vime.backends.megatron_utils.update_weight.update_weight_from_tensor" + sys.modules.pop(module_name, None) + module_path = REPO_ROOT / "vime" / "backends" / "megatron_utils" / "update_weight" / "update_weight_from_tensor.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + assert spec.loader is not None + spec.loader.exec_module(module) + return module, dist_state + + +def test_empty_colocated_bucket_still_participates_in_gather(monkeypatch): + module, dist_state = _load_update_weight_module(monkeypatch) + dist_state.gathered = lambda local: [local, local] + engine = _FakeEngine() + + refs, long_lived_tensor = module._send_to_colocated_engine( + [], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group=object(), + ) + + assert dist_state.local_object["names"] == [] + assert refs == [] + assert long_lived_tensor is None + assert engine.update_weights.calls == [] + + +def test_source_rank_sends_empty_colocated_bucket(monkeypatch): + module, dist_state = _load_update_weight_module(monkeypatch) + remote_info = { + "names": ["expert.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[4, 8]], + "tensor_sizes": [64], + "ipc_handles": {"gpu-1": ("remote",)}, + } + dist_state.gathered = lambda local: [local, remote_info] + engine = _FakeEngine() + + refs, long_lived_tensor = module._send_to_colocated_engine( + [], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group=object(), + ) + + assert refs == ["ref-1"] + assert long_lived_tensor is None + assert engine.update_weights.calls == [ + ( + ( + [ + { + "names": [], + "dtype_names": [], + "shapes": [], + "tensor_sizes": [], + "ipc_handles": {}, + }, + remote_info, + ], + ), + {}, + ) + ] + + +def test_source_rank_sends_different_expert_metadata_as_separate_updates(monkeypatch): + module, dist_state = _load_update_weight_module(monkeypatch) + local_info = { + "names": ["experts.0.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[4, 8]], + "tensor_sizes": [64], + "ipc_handles": {"gpu-0": ("local",)}, + } + remote_info = { + **local_info, + "names": ["experts.1.weight"], + "ipc_handles": {"gpu-1": ("remote",)}, + } + dist_state.gathered = lambda local: [local, remote_info] + engine = _FakeEngine() + + monkeypatch.setattr(module, "_build_packed_ipc_update_info", lambda _tensors: (local_info, "packed")) + refs, long_lived_tensor = module._send_to_colocated_engine( + [("experts.0.weight", object())], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group=object(), + ) + + assert refs == ["ref-1"] + assert long_lived_tensor == "packed" + assert engine.update_weights.calls == [(([local_info, remote_info],), {})] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_eval_config.py b/tests/test_eval_config.py new file mode 100644 index 000000000..41a5055d4 --- /dev/null +++ b/tests/test_eval_config.py @@ -0,0 +1,116 @@ +"""CPU unit tests for ``vime.utils.eval_config.build_eval_dataset_configs``. + +The documented contract (examples/eval_multi_task/README.md) is that +``eval.defaults`` "defines inference parameters shared by every dataset entry. +Override them inside an individual dataset block if needed." — i.e. resolution +is dataset entry > defaults > args, for every ``EvalDatasetConfig`` field. + +Historically only the fields listed in the two spec tables flowed through +``defaults``; everything else (``rm_type``, ``repetition_penalty``, +``app_service``, ...) was silently dropped, and a typo'd key in ``defaults`` +was silently accepted while the same typo in a dataset entry raised. These +tests pin the full contract, including the ``stop`` / ``stop_token_ids`` / +``min_new_tokens`` fields that lost their resolution in the #1005 refactor. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vime.utils.eval_config import build_eval_dataset_configs + + +NUM_GPUS = 0 + + +def _args(**overrides): + values = dict( + rollout_temperature=0.8, + rollout_top_p=1.0, + rollout_stop=[""], + rollout_stop_token_ids=[7], + eval_min_new_tokens=None, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +@pytest.mark.unit +def test_non_spec_defaults_reach_every_dataset(): + datasets = build_eval_dataset_configs( + _args(), + [{"name": "aime", "path": "/d/aime.jsonl"}, {"name": "gpqa", "path": "/d/gpqa.jsonl"}], + defaults={"rm_type": "deepscaler", "repetition_penalty": 1.05, "eval_task_timeout": 120}, + ) + + for dataset in datasets: + assert dataset.rm_type == "deepscaler" + assert dataset.repetition_penalty == 1.05 + assert dataset.eval_task_timeout == 120 + + +@pytest.mark.unit +def test_dataset_entry_overrides_default(): + datasets = build_eval_dataset_configs( + _args(), + [ + {"name": "aime", "path": "/d/aime.jsonl", "rm_type": "math", "temperature": 0.2}, + {"name": "gpqa", "path": "/d/gpqa.jsonl"}, + ], + defaults={"rm_type": "deepscaler", "temperature": 0.7}, + ) + + assert datasets[0].rm_type == "math" + assert datasets[0].temperature == 0.2 + assert datasets[1].rm_type == "deepscaler" + assert datasets[1].temperature == 0.7 + + +@pytest.mark.unit +def test_stop_fields_resolve_dataset_then_default_then_args(): + datasets = build_eval_dataset_configs( + _args(eval_min_new_tokens=4), + [ + {"name": "a", "path": "/d/a.jsonl", "stop": [""], "min_new_tokens": 8}, + {"name": "b", "path": "/d/b.jsonl"}, + {"name": "c", "path": "/d/c.jsonl"}, + ], + defaults={"stop_token_ids": [11, 12]}, + ) + + # dataset entry wins + assert datasets[0].stop == [""] + assert datasets[0].min_new_tokens == 8 + # eval.defaults fills in + assert datasets[1].stop_token_ids == [11, 12] + # args are the last fallback + assert datasets[1].stop == [""] + assert datasets[2].stop_token_ids == [11, 12] + assert datasets[2].min_new_tokens == 4 + + +@pytest.mark.unit +def test_unknown_default_key_raises(): + with pytest.raises(ValueError, match="temperture"): + build_eval_dataset_configs( + _args(), + [{"name": "aime", "path": "/d/aime.jsonl"}], + defaults={"temperture": 0.7}, + ) + + +@pytest.mark.unit +def test_spec_fields_still_fall_back_to_args(): + datasets = build_eval_dataset_configs( + _args(rollout_temperature=0.9), + [{"name": "aime", "path": "/d/aime.jsonl"}], + defaults={}, + ) + + assert datasets[0].temperature == 0.9 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_expert_routing.py b/tests/test_expert_routing.py new file mode 100644 index 000000000..1ea0d6882 --- /dev/null +++ b/tests/test_expert_routing.py @@ -0,0 +1,157 @@ +import sys +import types +from argparse import Namespace +from dataclasses import replace + +import pytest +import torch + +from vime.backends.megatron_utils.update_weight.expert_routing import ( + _build_expert_transfer_plan, + _can_route_experts, + _expert_transfer_size, + _ExpertParam, + _get_expert_target_ranks, + _get_vllm_moe_topology, +) +from vime.utils.types import ParamInfo + +NUM_GPUS = 0 + + +def _topology_args(**overrides): + values = { + "vllm_pp_size": 1, + "vllm_prefill_context_parallel_size": 1, + "vllm_dp_size": 1, + "vllm_enable_expert_parallel": False, + "vllm_enable_eplb": False, + "vllm_eplb_config": Namespace(num_redundant_experts=0), + "vllm_expert_placement_strategy": "linear", + "vllm_enable_elastic_ep": False, + } + values.update(overrides) + return Namespace(**values) + + +def test_vllm_moe_topology_derives_ep_from_data_parallelism(): + topology = _get_vllm_moe_topology( + _topology_args(vllm_dp_size=64, vllm_enable_expert_parallel=True), + engine_gpu_count=64, + ) + + assert (topology.tp_size, topology.pp_size, topology.pcp_size, topology.dp_size) == (1, 1, 1, 64) + assert topology.enable_expert_parallel is True + assert topology.ep_size == 64 + + +def test_vllm_moe_topology_includes_pcp_and_tp_in_ep_group(): + topology = _get_vllm_moe_topology( + _topology_args(vllm_prefill_context_parallel_size=2, vllm_enable_expert_parallel=True), + engine_gpu_count=4, + ) + + assert (topology.tp_size, topology.pcp_size, topology.dp_size) == (2, 2, 1) + assert topology.ep_size == 4 + + +def test_vllm_moe_topology_disables_expert_sharding_without_ep(): + topology = _get_vllm_moe_topology(_topology_args(), engine_gpu_count=4) + + assert topology.tp_size == 4 + assert topology.enable_expert_parallel is False + assert topology.ep_size == 1 + + +def test_expert_target_ranks_map_each_ep_shard_to_colocated_rank(): + assert _get_expert_target_ranks([4], [2], ep_size=4, world_size=8) == ((2,), (3,), (4,), (5,)) + + +def _can_route_with_args(monkeypatch, **overrides): + mpu = types.SimpleNamespace(get_expert_tensor_parallel_world_size=lambda: 1) + megatron = types.ModuleType("megatron") + megatron_core = types.ModuleType("megatron.core") + megatron.core = megatron_core + megatron_core.mpu = mpu + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", megatron_core) + args = _topology_args(vllm_enable_expert_parallel=True, **overrides) + topology = _get_vllm_moe_topology(args, engine_gpu_count=2) + return _can_route_experts(args, topology, engine_gpu_counts=[2]) + + +def test_vllm_rank_local_expert_routing_accepts_linear_static_ep(monkeypatch): + assert _can_route_with_args(monkeypatch) + + +@pytest.mark.parametrize( + "overrides", + [ + {"vllm_expert_placement_strategy": "round_robin"}, + {"vllm_enable_eplb": True}, + {"vllm_eplb_config": Namespace(num_redundant_experts=1)}, + {"vllm_enable_elastic_ep": True}, + ], + ids=["round-robin", "eplb", "redundant-experts", "elastic-ep"], +) +def test_vllm_rank_local_expert_routing_rejects_non_static_placement(monkeypatch, overrides): + assert not _can_route_with_args(monkeypatch, **overrides) + + +def _param(*, expert: int, projection: int, source_rank: int, target_rank: int, size: int) -> _ExpertParam: + info = ParamInfo( + name=f"module.module.decoder.layers.3.mlp.experts.linear_fc{projection}.weight{expert}", + dtype=torch.bfloat16, + shape=torch.Size([size // 2]), + attrs={}, + size=size, + src_rank=source_rank, + ) + return _ExpertParam( + info=info, + layer=3, + expert=expert, + target_ranks=(target_rank,), + ) + + +def test_transfer_plan_splits_same_rank_experts_at_expert_boundaries(): + params = [] + for expert, rank in ((0, 0), (1, 0), (2, 1), (3, 1)): + params.extend( + [ + _param(expert=expert, projection=1, source_rank=rank, target_rank=rank, size=60), + _param(expert=expert, projection=2, source_rank=rank, target_rank=rank, size=60), + ] + ) + + plan = _build_expert_transfer_plan(params, buffer_size=150) + + assert len(plan) == 1 + assert len(plan[0]) == 2 + transfers = [transfer for batch in plan[0] for transfer in batch] + assert len(transfers) == 4 + assert all(_expert_transfer_size(transfer) == 120 for transfer in transfers) + assert all(len({param.expert for param in transfer.params}) == 1 for transfer in transfers) + assert {(param.expert, param.info.name) for transfer in transfers for param in transfer.params} == { + (param.expert, param.info.name) for param in params + } + + +def test_transfer_plan_rejects_one_expert_larger_than_buffer(): + first = _param(expert=0, projection=1, source_rank=0, target_rank=0, size=80) + second = replace( + first, + info=replace( + first.info, + name="module.module.decoder.layers.3.mlp.experts.linear_fc2.weight0", + size=80, + ), + ) + + with pytest.raises(ValueError, match="exceeds update_weight_buffer_size"): + _build_expert_transfer_plan([first, second], buffer_size=150) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_external_vllm_engines.py b/tests/test_external_vllm_engines.py new file mode 100644 index 000000000..5e417ef9c --- /dev/null +++ b/tests/test_external_vllm_engines.py @@ -0,0 +1,242 @@ +import sys +import types +from argparse import Namespace +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.backends.vllm_utils.external import ( + ExternalEngineInfo, + apply_external_engine_info_to_args, + discover_external_engines, + get_server_info, + start_external_rollout_servers, +) +from vime.utils.http_utils import get_rollout_num_engines + +NUM_GPUS = 0 + + +class _Response: + def __init__(self, payload, status_code=200): + self.payload = payload + self.status_code = status_code + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def json(self): + return self.payload + + +def test_discover_external_engines_reads_server_info(monkeypatch): + def fake_get(url, timeout): + assert timeout == 30.0 + assert url == "http://host1:10090/server_info?config_format=json" + return _Response( + { + "tp_size": 4, + "pp_size": 2, + "pcp_size": 1, + "dp_size": 1, + "enable_expert_parallel": True, + "disaggregation_mode": "null", + } + ) + + monkeypatch.setattr("vime.backends.vllm_utils.external.requests.get", fake_get) + + infos = discover_external_engines(["host1:10090"]) + + assert len(infos) == 1 + info = infos[0] + assert info.url == "http://host1:10090" + assert info.host == "host1" + assert info.port == 10090 + assert info.worker_type == "regular" + assert info.num_gpus == 8 + assert info.server_info["tp_size"] == 4 + assert info.server_info["pp_size"] == 2 + assert info.server_info["dp_size"] == 1 + assert info.server_info["enable_expert_parallel"] is True + assert info.parallel_config == { + "tp_size": 4, + "pp_size": 2, + "pcp_size": 1, + "dp_size": 1, + "enable_expert_parallel": True, + "ep_size": 4, + } + + +def test_start_external_rollout_servers_exposes_parallel_configs(monkeypatch): + class FakeActor: + init = Namespace(remote=lambda **kwargs: kwargs) + + class FakeActorClass: + def options(self, **kwargs): + return self + + def remote(self, **kwargs): + return FakeActor() + + ray = types.ModuleType("ray") + ray.remote = lambda actor_class: FakeActorClass() + vllm_engine = types.ModuleType("vime.backends.vllm_utils.vllm_engine") + vllm_engine.VLLMEngine = object + ray_utils = types.ModuleType("vime.ray.utils") + ray_utils.add_default_ray_env_vars = lambda: {} + monkeypatch.setitem(sys.modules, "ray", ray) + monkeypatch.setitem(sys.modules, "vime.backends.vllm_utils.vllm_engine", vllm_engine) + monkeypatch.setitem(sys.modules, "vime.ray.utils", ray_utils) + + info = ExternalEngineInfo( + url="http://host1:10090", + host="host1", + port=10090, + worker_type="regular", + num_gpus=8, + server_info={ + "tp_size": 2, + "pp_size": 2, + "pcp_size": 1, + "dp_size": 2, + "enable_expert_parallel": True, + }, + ) + args = Namespace(rollout_external_engine_infos=[info.to_dict()]) + + servers, init_handles = start_external_rollout_servers( + args, + start_router=lambda *args, **kwargs: ("host1", 30000, None), + ) + + assert servers["default"].engine_parallel_configs == [ + { + "tp_size": 2, + "pp_size": 2, + "pcp_size": 1, + "dp_size": 2, + "enable_expert_parallel": True, + "ep_size": 4, + } + ] + assert len(init_handles) == 1 + + +def test_get_server_info_flattens_nested_vllm_transfer_configs(monkeypatch): + def fake_get(url, timeout): + assert url == "http://host1:10090/server_info?config_format=json" + assert timeout == 30.0 + return _Response( + { + "vllm_config": { + "parallel_config": {"tensor_parallel_size": 2}, + "model_config": { + "kv_transfer_config": { + "kv_connector": "NixlConnector", + "kv_role": "kv_producer", + }, + "weight_transfer_config": {"backend": "nccl"}, + }, + }, + "vllm_env": {"VLLM_NIXL_SIDE_CHANNEL_PORT": 12090}, + } + ) + + monkeypatch.setattr("vime.backends.vllm_utils.external.requests.get", fake_get) + server_info = get_server_info("http://host1:10090") + + assert server_info["tensor_parallel_size"] == 2 + assert server_info["kv_transfer_config"]["kv_role"] == "kv_producer" + assert server_info["weight_transfer_config"] == {"backend": "nccl"} + assert server_info["disaggregation_mode"] == "prefill" + assert server_info["disaggregation_bootstrap_port"] == 12090 + + +def test_apply_external_engine_info_handles_pd(monkeypatch): + payloads = { + "http://prefill:10090/server_info?config_format=json": { + "tp_size": 2, + "pp_size": 1, + "pcp_size": 1, + "dp_size": 1, + "enable_expert_parallel": False, + "disaggregation_mode": "prefill", + "disaggregation_bootstrap_port": 12090, + }, + "http://decode:10091/server_info?config_format=json": { + "tp_size": 4, + "pp_size": 1, + "pcp_size": 1, + "dp_size": 2, + "enable_expert_parallel": True, + "disaggregation_mode": "decode", + }, + } + + def fake_get(url, timeout): + return _Response(payloads[url]) + + monkeypatch.setattr("vime.backends.vllm_utils.external.requests.get", fake_get) + args = Namespace( + rollout_external=True, + rollout_external_engine_addrs=["prefill:10090", "decode:10091"], + rollout_num_gpus=None, + rollout_num_gpus_per_engine=1, + router_pd_disaggregation=False, + ) + + apply_external_engine_info_to_args(args) + + assert args.rollout_external is True + assert args.router_pd_disaggregation is False + assert args.rollout_num_gpus == 10 + assert args.rollout_num_engines == 2 + assert get_rollout_num_engines(args) == 2 + assert [info["worker_type"] for info in args.rollout_external_engine_infos] == ["prefill", "decode"] + assert [info["num_gpus"] for info in args.rollout_external_engine_infos] == [2, 8] + assert [info["server_info"]["dp_size"] for info in args.rollout_external_engine_infos] == [1, 2] + assert args.rollout_external_engine_infos[0]["disaggregation_bootstrap_port"] == 12090 + + +def test_apply_external_engine_info_preserves_router_pd_flag(monkeypatch): + def fake_get(url, timeout): + assert url == "http://regular:10090/server_info?config_format=json" + return _Response( + { + "tp_size": 2, + "pp_size": 1, + "disaggregation_mode": "null", + } + ) + + monkeypatch.setattr("vime.backends.vllm_utils.external.requests.get", fake_get) + args = Namespace( + rollout_external=True, + rollout_external_engine_addrs=["regular:10090"], + router_pd_disaggregation=True, + ) + + apply_external_engine_info_to_args(args) + + assert args.rollout_external is True + assert args.router_pd_disaggregation is True + assert args.rollout_num_gpus == 2 + assert args.rollout_num_engines == 1 + + +def test_apply_external_engine_info_requires_addrs(): + args = Namespace(rollout_external_engine_addrs=None) + + with pytest.raises(ValueError, match="rollout-external-engine-addrs"): + apply_external_engine_info_to_args(args) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_filter_long_prompt.py b/tests/test_filter_long_prompt.py new file mode 100644 index 000000000..f8129fa7c --- /dev/null +++ b/tests/test_filter_long_prompt.py @@ -0,0 +1,116 @@ +"""CPU unit tests for ``vime.utils.data.filter_long_prompt``. + +With a processor configured, the function scores text-only samples with a +batched tokenizer call and multimodal samples one at a time through the +processor. Splitting the work that way is a throughput optimization; it must not +change which samples survive, or the order they survive in. + +Order matters concretely: ``--rollout-shuffle`` defaults to False, so +``Dataset.samples`` is consumed in exactly the order this function returns. +Grouping the survivors by modality would make a mixed dataset train every +text-only prompt before any prompt carrying an image. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from vime.utils.data import filter_long_prompt +from vime.utils.types import Sample + + +NUM_GPUS = 0 + + +@pytest.fixture +def stub_processor_kwargs(monkeypatch): + """`vime.utils.processing_utils` pulls in transformers + PIL. + + The multimodal branch imports it lazily, and only `build_processor_kwargs` is + needed here, so stub the module rather than requiring transformers on the + CPU image. + """ + module = types.ModuleType("vime.utils.processing_utils") + module.build_processor_kwargs = lambda multimodal_inputs: {"images": None} + monkeypatch.setitem(sys.modules, "vime.utils.processing_utils", module) + + +class _Tokenizer: + """Batched tokenizer stand-in: prompt "pN:len" tokenizes to `len` ids.""" + + def __call__(self, prompts, add_special_tokens=False): + return {"input_ids": [[0] * _encoded_length(p) for p in prompts]} + + +class _Processor: + """Per-sample processor stand-in, same length convention.""" + + def __call__(self, text=None, **kwargs): + return {"input_ids": [[0] * _encoded_length(text)]} + + +def _encoded_length(prompt: str) -> int: + return int(prompt.split(":")[1]) + + +def _make_samples(specs): + """specs: list of (is_multimodal, encoded_length).""" + samples = [] + for i, (is_multimodal, length) in enumerate(specs): + sample = Sample(prompt=f"p{i}:{length}") + sample.multimodal_inputs = {"images": ["img"]} if is_multimodal else None + samples.append(sample) + return samples + + +@pytest.mark.unit +def test_preserves_order_when_nothing_is_filtered(stub_processor_kwargs): + # Alternating modality, every prompt well under the limit. + samples = _make_samples([(i % 2 == 0, 5) for i in range(6)]) + + kept = filter_long_prompt(samples, _Tokenizer(), _Processor(), max_length=100) + + assert [s.prompt for s in kept] == [s.prompt for s in samples] + + +@pytest.mark.unit +def test_preserves_order_when_some_are_filtered(stub_processor_kwargs): + specs = [ + (True, 5), # p0 multimodal, keep + (False, 500), # p1 text-only, drop + (False, 5), # p2 text-only, keep + (True, 500), # p3 multimodal, drop + (False, 5), # p4 text-only, keep + (True, 5), # p5 multimodal, keep + ] + samples = _make_samples(specs) + + kept = filter_long_prompt(samples, _Tokenizer(), _Processor(), max_length=100) + + assert [s.prompt for s in kept] == ["p0:5", "p2:5", "p4:5", "p5:5"] + + +@pytest.mark.unit +@pytest.mark.parametrize("all_multimodal", [True, False]) +def test_single_modality_batches_are_unchanged(stub_processor_kwargs, all_multimodal): + samples = _make_samples([(all_multimodal, 5)] * 4) + + kept = filter_long_prompt(samples, _Tokenizer(), _Processor(), max_length=100) + + assert [s.prompt for s in kept] == [s.prompt for s in samples] + + +@pytest.mark.unit +def test_no_processor_path_still_preserves_order(): + samples = _make_samples([(False, 5), (False, 500), (False, 5)]) + + kept = filter_long_prompt(samples, _Tokenizer(), None, max_length=100) + + assert [s.prompt for s in kept] == ["p0:5", "p2:5"] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_full_disk_weight_update.py b/tests/test_full_disk_weight_update.py new file mode 100644 index 000000000..3eacae45b --- /dev/null +++ b/tests/test_full_disk_weight_update.py @@ -0,0 +1,139 @@ +"""E2E smoke test for full checkpoint weight updates through disk. + +Runs a tiny Qwen3.5-0.8B job where each weight sync writes a complete HF +checkpoint and rollout engines reload it through ``update_weights_from_disk``. +""" + +import os +import tempfile +from pathlib import Path + +import vime.utils.external_utils.command_utils as U + + +MODEL_NAME = "Qwen3.5-0.8B" +MODEL_TYPE = "qwen3.5-0.8B" +NUM_GPUS = 4 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) + + +def execute(): + with tempfile.TemporaryDirectory(prefix="vime_full_disk_weight_update_") as disk_dir: + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load {TORCH_DIST_CKPT} " + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 1 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 1024 " + "--rollout-temperature 0.8 " + "--over-sampling-batch-size 8 " + "--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " + "--global-batch-size 16 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.01 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 1 " + "--rollout-num-gpus 3 " + "--vllm-gpu-memory-utilization 0.7 " + "--vllm-max-cudagraph-capture-size 32 " + ) + + disk_update_args = ( + "--update-weight-mode full " + "--update-weight-transport disk " + f"--update-weight-disk-dir {disk_dir} " + "--update-weight-disk-keep-files " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--loss-mask-type qwen3_5 " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 1 " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{vllm_args} " + f"{disk_update_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + checkpoint_dirs = sorted(Path(disk_dir).glob("weight_v*")) + assert checkpoint_dirs, f"No disk checkpoint directories were written under {disk_dir}" + assert any((path / "model.safetensors.index.json").exists() for path in checkpoint_dirs) + assert any(list(path.glob("*.safetensors")) for path in checkpoint_dirs) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_fully_async_rollout.py b/tests/test_fully_async_rollout.py new file mode 100644 index 000000000..bdf320f02 --- /dev/null +++ b/tests/test_fully_async_rollout.py @@ -0,0 +1,175 @@ +"""CPU unit tests for the fully-async rollout worker's queue contract. + +The module docstring of ``vime.rollout.fully_async_rollout`` promises that the +worker's output queue "stays warm" across ``generate_rollout`` calls: each call +takes ``rollout_batch_size`` completed groups and leaves the rest queued. + +Three behaviours are pinned here: + + 1. ``_generate_rollout_async`` consumes exactly ``rollout_batch_size`` groups + and leaves the surplus in the queue. (It used to drain the whole queue and + slice — throwing away fully generated, reward-scored groups whose prompts + had already been consumed from the data buffer.) + 2. The task done-callback never blocks. It runs on the event-loop thread, so + a bounded queue that filled up would freeze every in-flight generation. + 3. Backpressure exists anyway: ``_loop`` stops pulling new prompts while a + full pool of completed groups is already waiting to be consumed. +""" + +from __future__ import annotations + +import asyncio +import sys +import threading +import time +import types +from collections import deque +from types import SimpleNamespace + +# ``fully_async_rollout`` imports ``vllm_rollout``, which needs vllm_router +# and (transitively) transformers — both deliberately absent from the CPU CI +# env. The tests below never dial a server or touch a tokenizer, so stub the +# imports, same as tests/test_agent/test_agent_rollout_cpu.py. +if "vllm_router" not in sys.modules: + _router_stub = types.ModuleType("vllm_router") + _router_stub.__version__ = "0.2.3" + sys.modules["vllm_router"] = _router_stub +if "transformers" not in sys.modules: + _tf_stub = types.ModuleType("transformers") + for _name in ("AutoProcessor", "AutoTokenizer", "PreTrainedTokenizerBase", "ProcessorMixin"): + setattr(_tf_stub, _name, type(_name, (), {})) + sys.modules["transformers"] = _tf_stub + +import pytest + +import vime.rollout.fully_async_rollout as fa +from vime.utils.types import Sample + + +NUM_GPUS = 0 + + +class _FakeGenerateState: + def __init__(self, args): + self.sampling_params = {} + + +class _FakeDataBuffer: + """Finite fuel: one group per ``get_samples`` call until exhausted.""" + + def __init__(self, groups): + self._groups = deque(groups) + self.requeued = [] + + def get_samples(self, n): + assert n == 1 + if not self._groups: + return [] + return [self._groups.popleft()] + + def add_samples(self, groups): + self.requeued.extend(groups) + + +def _make_group(index: int) -> list[Sample]: + sample = Sample(index=index, prompt=f"p{index}") + sample.status = Sample.Status.COMPLETED + return [sample] + + +def _make_worker(monkeypatch, data_buffer=None, concurrency=4) -> fa.AsyncRolloutWorker: + monkeypatch.setattr(fa, "GenerateState", _FakeGenerateState) + args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4) + return fa.AsyncRolloutWorker(args, data_buffer or _FakeDataBuffer([]), concurrency=concurrency) + + +@pytest.mark.unit +def test_rollout_takes_target_groups_and_leaves_surplus_queued(monkeypatch): + worker = _make_worker(monkeypatch) + for gid in range(10): + worker.output_queue.put((gid, _make_group(gid))) + monkeypatch.setattr(fa, "_get_global_worker", lambda args, data_buffer: worker) + + args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4) + out = asyncio.run(fa._generate_rollout_async(args, rollout_id=0, data_buffer=None)) + + assert len(out) == 4 + # FIFO: the oldest four groups ship first. + assert [group[0].index for group in out] == [0, 1, 2, 3] + # The other six are still queued for the next rollout, not thrown away. + assert worker.queue_size() == 6 + assert [gid for gid, _ in worker.get_completed_groups()] == [4, 5, 6, 7, 8, 9] + + +@pytest.mark.unit +def test_get_completed_groups_limit(monkeypatch): + worker = _make_worker(monkeypatch) + for gid in range(5): + worker.output_queue.put((gid, _make_group(gid))) + + assert [gid for gid, _ in worker.get_completed_groups(limit=2)] == [0, 1] + assert [gid for gid, _ in worker.get_completed_groups()] == [2, 3, 4] + assert worker.get_completed_groups(limit=3) == [] + + +@pytest.mark.unit +def test_done_callback_never_blocks_event_loop_thread(monkeypatch): + """The callback runs on the loop thread; blocking there freezes every + in-flight generation. Push more results than the old bounded-queue cap + (1000) through it and require completion.""" + worker = _make_worker(monkeypatch) + + class _DoneTask: + def __init__(self, gid): + self._result = _make_group(gid) + + def result(self): + return self._result + + def _push_all(): + for gid in range(1001): + worker._make_done_cb(gid)(_DoneTask(gid)) + + pusher = threading.Thread(target=_push_all, daemon=True) + pusher.start() + pusher.join(timeout=30) + + assert not pusher.is_alive(), "done-callback blocked on a full output queue" + assert worker.queue_size() == 1001 + + +@pytest.mark.unit +def test_loop_backpressure_stops_topping_up_when_queue_is_full(monkeypatch): + """With instantly-completing generations and plenty of fuel, the queue must + plateau around ``concurrency`` instead of absorbing the whole dataset.""" + concurrency = 3 + fuel = 60 + data_buffer = _FakeDataBuffer([_make_group(i) for i in range(fuel)]) + + async def _instant_generate(args, group, sampling_params, evaluation): + return group + + monkeypatch.setattr(fa, "generate_and_rm_group", _instant_generate) + worker = _make_worker(monkeypatch, data_buffer=data_buffer, concurrency=concurrency) + worker.poll_interval = 0.01 + + worker.start() + try: + # Give the loop ample iterations to overshoot if it is going to. + deadline = time.time() + 3.0 + max_seen = 0 + while time.time() < deadline: + max_seen = max(max_seen, worker.queue_size()) + if max_seen > 2 * concurrency: + break + time.sleep(0.02) + finally: + worker.stop() + + # In-flight tasks may still land after the gate check, so allow one pool + # beyond the gate — but nothing near the unthrottled fuel size. + assert 0 < max_seen <= 2 * concurrency, f"queue grew to {max_seen} with concurrency={concurrency}" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_geo3k_vlm_multi_turn_e2e.py b/tests/test_geo3k_vlm_multi_turn_e2e.py new file mode 100644 index 000000000..55e6065da --- /dev/null +++ b/tests/test_geo3k_vlm_multi_turn_e2e.py @@ -0,0 +1,281 @@ +"""Single-GPU end-to-end regression for Geo3K Issue #331.""" + +from __future__ import annotations + +import asyncio +import os +import shlex +import subprocess +import sys +import time +from argparse import Namespace +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.error import URLError +from urllib.parse import urlsplit +from urllib.request import urlopen + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from vllm.utils.system_utils import kill_process_tree + +import vime.utils.external_utils.command_utils as U +from vime.utils.http_utils import is_port_available +from vime.utils.misc import load_function + +NUM_GPUS = 1 +MODEL_NAME = "Qwen3-VL-2B-Instruct" +MODEL_REVISION = "89644892e4d85e24eaac8bacfd4f463576704203" +TEST_ROOT = Path(os.environ.get("VIME_TEST_ROOT", "/root")) +MODEL_PATH = TEST_ROOT / "models" / MODEL_NAME +DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed" +DATASET_REVISION = "53ff8fbb1f9758b9efbc005e9487e1fdd874364a" +DATA_ROOT = TEST_ROOT / "datasets/geo3k_imgurl_processed" +TRAIN_DATA_PATH = DATA_ROOT / "train.parquet" +CUSTOM_GENERATE_PATH = "examples.geo3k_vlm_multi_turn.rollout.generate" +HOST = "127.0.0.1" +WORKER_PORT = 10090 +ROUTER_PORT = 30000 +PROMETHEUS_PORT = 29001 +WORKER_URL = f"http://{HOST}:{WORKER_PORT}" +ROUTER_URL = f"http://{HOST}:{ROUTER_PORT}" +RENDER_ENDPOINT = "/v1/chat/completions/render" +GENERATE_ENDPOINT = "/inference/v1/generate" +SEED = 331 +ROW_INDEX = 0 +MAX_NEW_TOKENS = 1024 +MAX_TURNS = 2 +TEMPERATURE = 0.0 +TOP_P = 1.0 +TOP_K = -1 +GPU_MEMORY_UTILIZATION = 0.70 +MAX_MODEL_LEN = 4096 +MAX_NUM_SEQS = 1 +ROUTER_REQUEST_TIMEOUT_S = 600 +HTTP_MAX_RETRIES = 60 +SERVICE_START_TIMEOUT_S = 180 +SERVICE_POLL_INTERVAL_S = 2 +PROCESS_STOP_TIMEOUT_S = 30 +WORKER_LOG = Path("/tmp/geo3k-vllm-worker.log") +ROUTER_LOG = Path("/tmp/geo3k-vllm-router.log") +WORKER_COMMAND = shlex.split( + f"vllm serve {MODEL_PATH} --host {HOST} --port {WORKER_PORT} --dtype bfloat16 " + f"--tensor-parallel-size {NUM_GPUS} --gpu-memory-utilization {GPU_MEMORY_UTILIZATION} " + f"--max-model-len {MAX_MODEL_LEN} --max-num-seqs {MAX_NUM_SEQS} --enforce-eager " + "--generation-config vllm --logprobs-mode processed_logprobs" +) +ROUTER_COMMAND = shlex.split( + f"vllm-router --host {HOST} --port {ROUTER_PORT} --worker-urls {WORKER_URL} " + f"--policy consistent_hash --prometheus-port {PROMETHEUS_PORT} --prometheus-host {HOST} " + f"--request-timeout-secs {ROUTER_REQUEST_TIMEOUT_S}" +) + + +@dataclass(frozen=True) +class RequestEvent: + endpoint: str + request_token_ids: tuple[int, ...] + response_token_ids: tuple[int, ...] + has_features: bool + + +class TwoTurnEnvironment: + def reset(self) -> None: + pass + + def step(self, _response: str) -> tuple[str, bool, dict]: + return "Continue the reasoning.", False, {} + + def format_observation(self, observation: str) -> dict[str, str]: + return {"role": "user", "content": observation} + + def close(self) -> None: + pass + + +def build_env(*, sample: Any, args: Namespace) -> TwoTurnEnvironment: + del sample, args + return TwoTurnEnvironment() + + +def prepare() -> None: + MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) + DATA_ROOT.parent.mkdir(parents=True, exist_ok=True) + U.exec_command(f"hf download Qwen/{MODEL_NAME} --revision {MODEL_REVISION} --local-dir {MODEL_PATH}") + U.exec_command( + f"hf download --repo-type dataset {DATASET_NAME} --revision {DATASET_REVISION} --local-dir {DATA_ROOT}" + ) + if not TRAIN_DATA_PATH.is_file(): + raise FileNotFoundError(f"Dataset not found at {TRAIN_DATA_PATH}") + + +def _wait_for_health(process: subprocess.Popen, health_url: str, log_path: Path) -> None: + deadline = time.monotonic() + SERVICE_START_TIMEOUT_S + while time.monotonic() < deadline: + if process.poll() is not None: + break + try: + with urlopen(f"{health_url}/health", timeout=SERVICE_POLL_INTERVAL_S): + return + except URLError: + time.sleep(SERVICE_POLL_INTERVAL_S) + logs = log_path.read_text(encoding="utf-8", errors="replace") + raise RuntimeError(f"Service failed to become healthy at {health_url}:\n{logs}") + + +def _stop_process(process: subprocess.Popen | None) -> None: + if process is None or process.poll() is not None: + return + kill_process_tree(process.pid) + process.wait(timeout=PROCESS_STOP_TIMEOUT_S) + + +def _start_process( + command: list[str], + health_url: str, + *, + log_path: Path, + env: dict[str, str] | None = None, +) -> subprocess.Popen: + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen(command, stdout=log_file, stderr=subprocess.STDOUT, env=env) + try: + _wait_for_health(process, health_url, log_path) + except BaseException: + _stop_process(process) + raise + return process + + +@contextmanager +def _vllm_services() -> Iterator[None]: + ports = (WORKER_PORT, ROUTER_PORT, PROMETHEUS_PORT) + if unavailable := [port for port in ports if not is_port_available(port)]: + raise RuntimeError(f"Required test ports are already in use: {unavailable}") + worker_env = os.environ.copy() + worker_env.update(VLLM_SERVER_DEV_MODE="1", VLLM_BATCH_INVARIANT="1") + worker = _start_process(WORKER_COMMAND, WORKER_URL, log_path=WORKER_LOG, env=worker_env) + router: subprocess.Popen | None = None + try: + router = _start_process(ROUTER_COMMAND, ROUTER_URL, log_path=ROUTER_LOG) + yield + finally: + try: + _stop_process(router) + finally: + _stop_process(worker) + + +@contextmanager +def _trace_http_requests() -> Iterator[list[RequestEvent]]: + from vime.utils import http_utils + + original_post = http_utils.post + events: list[RequestEvent] = [] + + async def traced_post(url, body, max_retries=HTTP_MAX_RETRIES, *, headers=None): + output = await original_post(url, body, max_retries=max_retries, headers=headers) + choices = output.get("choices") if isinstance(output, dict) else None + choice = choices[0] if isinstance(choices, list) and choices else {} + events.append( + RequestEvent( + endpoint=urlsplit(url).path, + request_token_ids=tuple(int(token) for token in body.get("token_ids") or ()), + response_token_ids=tuple(int(token) for token in choice.get("token_ids") or ()), + has_features=body.get("features") is not None, + ) + ) + return output + + http_utils.post = traced_post + try: + yield events + finally: + http_utils.post = original_post + + +def _load_sample() -> Any: + from vime.utils.data import Dataset + from vime.utils.processing_utils import load_processor, load_tokenizer + + data_slice = f"{TRAIN_DATA_PATH}@[{ROW_INDEX}:{ROW_INDEX + 1}]" + tokenizer = load_tokenizer(str(MODEL_PATH), trust_remote_code=True) + processor = load_processor(str(MODEL_PATH), trust_remote_code=True) + dataset = Dataset( + data_slice, + tokenizer, + processor, + None, + prompt_key="problem", + multimodal_keys={"image": "images"}, + ) + return dataset[0] + + +async def _run_rollout() -> tuple[Any, list[RequestEvent]]: + from vime.utils import http_utils + + args = Namespace( + partial_rollout=False, + max_turns=MAX_TURNS, + rollout_interaction_env_path=__name__, + vllm_router_ip=HOST, + vllm_router_port=ROUTER_PORT, + router_policy="consistent_hash", + hf_checkpoint=str(MODEL_PATH), + rollout_max_context_len=None, + use_rollout_routing_replay=False, + vllm_server_concurrency=NUM_GPUS, + rollout_num_engines=NUM_GPUS, + use_distributed_post=False, + rollout_temperature=TEMPERATURE, + rollout_top_p=TOP_P, + rollout_top_k=TOP_K, + rollout_max_response_len=MAX_NEW_TOKENS, + rollout_stop=None, + rollout_stop_token_ids=None, + rollout_skip_special_tokens=False, + vllm_dp_size=NUM_GPUS, + ) + sample = _load_sample() + http_utils.init_http_client(args) + sampling = { + "max_new_tokens": MAX_NEW_TOKENS, + "temperature": TEMPERATURE, + "top_p": TOP_P, + "top_k": TOP_K, + "skip_special_tokens": False, + "seed": SEED, + } + with _trace_http_requests() as events: + generate = load_function(CUSTOM_GENERATE_PATH) + result = await generate(args, sample, sampling) + return result, events + + +def _validate_result(sample: Any, events: list[RequestEvent]) -> None: + assert [event.endpoint for event in events] == [RENDER_ENDPOINT, GENERATE_ENDPOINT, GENERATE_ENDPOINT] + first, second = events[1:] + assert first.request_token_ids and first.response_token_ids + generated_start = len(first.request_token_ids) + generated_end = generated_start + len(first.response_token_ids) + assert second.request_token_ids[:generated_start] == first.request_token_ids + assert second.request_token_ids[generated_start:generated_end] == first.response_token_ids + assert len(second.request_token_ids) > generated_end + assert sample.tokens == list(second.request_token_ids + second.response_token_ids) + assert first.has_features and second.has_features + assert sample.response_length == len(sample.loss_mask) == len(sample.rollout_log_probs) + + +def execute() -> None: + with _vllm_services(): + sample, events = asyncio.run(_run_rollout()) + _validate_result(sample, events) + + +if __name__ == "__main__": + prepare() + execute() diff --git a/tests/test_glm4.7_30B_A3B_npu.py b/tests/test_glm4.7_30B_A3B_npu.py index c46d87d1c..27f054fe2 100644 --- a/tests/test_glm4.7_30B_A3B_npu.py +++ b/tests/test_glm4.7_30B_A3B_npu.py @@ -24,14 +24,8 @@ def execute(): model_dir = shlex.quote(MODEL_DIR) prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") - # NPU skips torch_dist conversion; HF weights load directly via bridge mode. - checkpoint_args = ( - f"--hf-checkpoint {model_dir} " - f"--load {model_dir} " - f"--ref-load {model_dir} " - "--megatron-to-hf-mode bridge " - "--no-load-optim " - ) + # Load HF weights through the native loader, without torch_dist conversion. + checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --ref-load {model_dir} --no-load-optim " # Smoke-scaled rollout (num-rollout/batch/n-samples trimmed like test_qwen3_30B_A3B_npu). rollout_args = ( @@ -44,7 +38,7 @@ def execute(): "--num-rollout 2 " "--rollout-batch-size 4 " "--n-samples-per-prompt 4 " - "--rollout-max-response-len 1024 " + "--rollout-max-response-len 128 " "--rollout-temperature 1 " "--global-batch-size 16 " "--balance-data " @@ -89,19 +83,19 @@ def execute(): "--use-precision-aware-optimizer " ) - # MTP speculative decoding under cudagraph - exercises the GLM-4.7 MTP - # drafter's graph-friendly forward (patched via docker/npu_patch/vllm.patch). - mtp_args = "--mtp-num-layers 1 " "--enable-mtp-training " "--mtp-loss-scaling-factor 0.2 " - vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " "--rollout-num-gpus-per-engine 4 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-enable-expert-parallel " "--vllm-cudagraph-capture-sizes 1 2 4 8 " '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":1}\' ' ) + mtp_args = "--mtp-num-layers 1 --enable-mtp-training --mtp-loss-scaling-factor 0.2 " model_args = ( + # GLM-4.7-Flash has no HF rope_scaling; MLA otherwise defaults to YaRN. + "--rope-type rope " "--attention-dropout 0.0 " "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " diff --git a/tests/test_glm4.7_30B_A3B_pd_mooncake.py b/tests/test_glm4.7_30B_A3B_pd_mooncake.py index 817e409ff..9247c691a 100644 --- a/tests/test_glm4.7_30B_A3B_pd_mooncake.py +++ b/tests/test_glm4.7_30B_A3B_pd_mooncake.py @@ -68,7 +68,8 @@ def execute(): "--rollout-batch-size 4 " "--n-samples-per-prompt 2 " "--rollout-max-response-len 512 " - "--rollout-temperature 0.8 " + "--rollout-temperature 1.0 " + "--rollout-top-p 1.0 " "--global-batch-size 8 " ) optimizer_args = ( @@ -112,8 +113,8 @@ def execute(): "--vllm-enable-expert-parallel " "--vllm-gpu-memory-utilization 0.45 " "--vllm-max-num-seqs 16 " - "--vllm-max-cudagraph-capture-size 8 " - '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":3}\' ' + "--vllm-max-cudagraph-capture-size 40 " + '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":4}\' ' "--router-request-timeout-secs 1200 " f"--vllm-config {vllm_config} " ) diff --git a/tests/test_glm52_6layer_deterministic_e2e.py b/tests/test_glm52_6layer_deterministic_e2e.py new file mode 100644 index 000000000..5878e6634 --- /dev/null +++ b/tests/test_glm52_6layer_deterministic_e2e.py @@ -0,0 +1,13 @@ +"""GLM-5.2 deterministic train/rollout alignment gate.""" + +NUM_GPUS = 8 + + +def run_gate(*, layerwise_zero: bool = False, rollout_max_response_len: int = 4096) -> None: + del layerwise_zero, rollout_max_response_len + # vLLM sparse MLA does not support batch-invariant inference. + raise RuntimeError("GLM-5.2 deterministic alignment is temporarily unsupported with vLLM") + + +def test_glm52_6layer_deterministic_train_rollout_alignment(): + run_gate() diff --git a/tests/test_glm52_layerwise_comparison.py b/tests/test_glm52_layerwise_comparison.py new file mode 100644 index 000000000..1e68e585d --- /dev/null +++ b/tests/test_glm52_layerwise_comparison.py @@ -0,0 +1,89 @@ +import pytest +import torch + +from glm52_layerwise_comparator import ( + TrainSequence, + _vllm_layer_token_rows, + compare_layer_outputs, + load_train_sequences, + map_requests_to_train_sequences, +) + +NUM_GPUS = 0 + + +def test_train_sequences_use_adjacent_cumulative_offsets(tmp_path): + dump_dir = tmp_path / "megatron" + dump_file = dump_dir / "rank00000" / "actor_Pass00000.pt" + dump_file.parent.mkdir(parents=True) + torch.save( + { + "input_ids": torch.tensor([7, 8, 9]), + "cu_seqlens": torch.tensor([0, 2, 3]), + "layers": {0: torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)}, + }, + dump_file, + ) + + sequences = load_train_sequences(dump_dir, {0}) + + assert [sequence.tokens.tolist() for sequence in sequences] == [[7, 8], [9]] + assert [sequence.layers[0].flatten().tolist() for sequence in sequences] == [[1.0, 2.0], [3.0]] + + +def test_health_check_requests_are_excluded_from_sequence_mapping(tmp_path): + dump_file = tmp_path / "Chunk00000.pt" + torch.save( + { + "model.forward_batch_info.input_ids": torch.tensor([99, 7, 8]), + "model.forward_batch_info.positions": torch.tensor([0, 0, 1]), + "model.forward_batch_info.rids": ["HEALTH_CHECK_probe", "rollout-0"], + "model.forward_batch_info.extend_seq_lens": torch.tensor([1, 2]), + }, + dump_file, + ) + train_sequences = [TrainSequence(tokens=torch.tensor([7, 8]), layers={}, source="test")] + + assert map_requests_to_train_sequences([dump_file], train_sequences) == {"rollout-0": 0} + + +def test_vllm_layer_output_reconstructs_the_visible_residual_sum(): + delta = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16) + residual = torch.tensor([[3.0, 4.0]], dtype=torch.bfloat16) + + output = _vllm_layer_token_rows([delta, residual], 1, "layer0") + + torch.testing.assert_close(output, torch.tensor([[4.0, 6.0]], dtype=torch.bfloat16)) + + +def test_layerwise_comparison_excludes_terminal_causal_state(tmp_path): + dump_file = tmp_path / "Chunk00000.pt" + torch.save( + { + "model.forward_batch_info.input_ids": torch.tensor([7, 8, 9]), + "model.forward_batch_info.positions": torch.tensor([0, 1, 2]), + "model.forward_batch_info.rids": ["rollout-0"], + "model.forward_batch_info.extend_seq_lens": torch.tensor([3]), + "model.layers.0": [ + torch.tensor([[1.0], [2.0], [100.0]], dtype=torch.bfloat16), + torch.zeros(3, 1, dtype=torch.bfloat16), + ], + }, + dump_file, + ) + train_sequences = [ + TrainSequence( + tokens=torch.tensor([7, 8, 9]), + layers={0: torch.tensor([[1.0], [2.0], [-100.0]], dtype=torch.bfloat16)}, + source="test", + ) + ] + + stats = compare_layer_outputs([dump_file], train_sequences, {"rollout-0": 0}, {0}) + + assert stats[0]["tokens"] == 2 + assert stats[0]["max_abs"] == 0.0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_glm52_layerwise_zero_e2e.py b/tests/test_glm52_layerwise_zero_e2e.py new file mode 100644 index 000000000..93a7ffc57 --- /dev/null +++ b/tests/test_glm52_layerwise_zero_e2e.py @@ -0,0 +1,18 @@ +"""Exact per-layer GLM-5 train/rollout alignment gate.""" + +import pytest + +from test_glm52_6layer_deterministic_e2e import run_gate + +NUM_GPUS = 8 + + +def test_glm52_first_six_layers_match_exactly(): + # Keep this diagnostic gate short: the main 4096-token test owns the + # realistic full-parameter training/logprob bound, while this run records + # every visible decoder-layer boundary and requires bitwise equality. + run_gate(layerwise_zero=True, rollout_max_response_len=32) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-s", "-rs"])) diff --git a/tests/test_glm5_indexer_q_norm.py b/tests/test_glm5_indexer_q_norm.py new file mode 100644 index 000000000..114633804 --- /dev/null +++ b/tests/test_glm5_indexer_q_norm.py @@ -0,0 +1,81 @@ +from types import SimpleNamespace + +import pytest +import torch + +from vime_plugins.models.glm5.glm5 import DSAMLASelfAttention, IdentityOp + +NUM_GPUS = 0 + + +class _FusedQUp(torch.nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.layer_norm_weight = torch.nn.Parameter(weight.clone()) + + +def _make_attention(weight: torch.Tensor, *, zero_centered_gamma: bool = False) -> DSAMLASelfAttention: + attention = DSAMLASelfAttention.__new__(DSAMLASelfAttention) + torch.nn.Module.__init__(attention) + attention.config = SimpleNamespace( + normalization="RMSNorm", + layernorm_epsilon=1.0e-5, + layernorm_zero_centered_gamma=zero_centered_gamma, + ) + attention.q_layernorm = IdentityOp() + attention.linear_q_up_proj = _FusedQUp(weight) + return attention + + +def test_glm5_indexer_uses_fused_q_rmsnorm_without_upstream_gradients(): + raw_q = torch.tensor( + [[1.0, -2.0, 3.0, -4.0], [0.25, 0.5, -0.75, 1.0]], + dtype=torch.bfloat16, + requires_grad=True, + ) + weight = torch.tensor([0.5, 1.0, 1.5, 2.0], dtype=torch.bfloat16) + attention = _make_attention(weight) + + actual = attention._get_indexer_q_input(raw_q) + expected = torch.nn.functional.rms_norm( + raw_q.detach().float(), + normalized_shape=(raw_q.shape[-1],), + weight=weight.float(), + eps=attention.config.layernorm_epsilon, + ).to(raw_q.dtype) + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + assert not torch.equal(actual, raw_q.detach()) + assert not actual.requires_grad + + wq_b = torch.nn.Linear(raw_q.shape[-1], 3, bias=False, dtype=torch.bfloat16) + wq_b(actual).float().sum().backward() + assert wq_b.weight.grad is not None + assert raw_q.grad is None + assert attention.linear_q_up_proj.layer_norm_weight.grad is None + + +def test_glm5_indexer_q_rmsnorm_supports_zero_centered_gamma_and_unfused_norm(): + raw_q = torch.tensor([[1.0, 2.0, 4.0, 8.0]], dtype=torch.bfloat16) + stored_weight = torch.tensor([-0.5, 0.0, 0.5, 1.0], dtype=torch.bfloat16) + attention = _make_attention(stored_weight, zero_centered_gamma=True) + + actual = attention._get_indexer_q_input(raw_q) + expected = torch.nn.functional.rms_norm( + raw_q.float(), + normalized_shape=(raw_q.shape[-1],), + weight=stored_weight.float() + 1.0, + eps=attention.config.layernorm_epsilon, + ).to(raw_q.dtype) + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + del attention.linear_q_up_proj.layer_norm_weight + attention.q_layernorm = torch.nn.RMSNorm(raw_q.shape[-1], eps=attention.config.layernorm_epsilon).to( + dtype=torch.bfloat16 + ) + unfused = attention._get_indexer_q_input(raw_q) + torch.testing.assert_close(unfused, attention.q_layernorm(raw_q), rtol=0.0, atol=0.0) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_glm5_indexer_short_context.py b/tests/test_glm5_indexer_short_context.py new file mode 100644 index 000000000..9822bd718 --- /dev/null +++ b/tests/test_glm5_indexer_short_context.py @@ -0,0 +1,55 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + +NUM_GPUS = 0 + + +def _load_indexer_module(): + module_name = "vime_plugins.models.glm5.ops.indexer_short_context_test" + for package_name in ( + "vime_plugins", + "vime_plugins.models", + "vime_plugins.models.glm5", + "vime_plugins.models.glm5.ops", + ): + package = sys.modules.setdefault(package_name, types.ModuleType(package_name)) + package.__path__ = [] + + for dependency in ("tilelang_indexer_bwd", "tilelang_indexer_fwd"): + dependency_name = f"vime_plugins.models.glm5.ops.{dependency}" + dependency_module = types.ModuleType(dependency_name) + setattr( + dependency_module, + "indexer_bwd_interface" if dependency.endswith("bwd") else "indexer_fwd_interface", + None, + ) + sys.modules[dependency_name] = dependency_module + + source = Path(__file__).parents[1] / "vime_plugins/models/glm5/ops/indexer.py" + spec = importlib.util.spec_from_file_location(module_name, source) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def test_short_context_topk_is_padded_with_invalid_routes(): + indexer = _load_indexer_module() + logits = torch.tensor([[0.5, float("-inf"), 0.25]]) + + scores, indices = indexer.pytorch_topk_with_invalid_padding(logits, topk=5) + + assert scores.shape == (1, 5) + assert indices.dtype == torch.int32 + assert sorted(indices[0][indices[0] >= 0].tolist()) == [0, 2] + assert indices[0].tolist().count(-1) == 3 + assert torch.isneginf(scores[0, 2:]).all() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_gspo.sh b/tests/test_gspo.sh deleted file mode 100644 index dd767abba..000000000 --- a/tests/test_gspo.sh +++ /dev/null @@ -1,79 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONUNBUFFERED=1 - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-0.6B -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type deepscaler - --num-rollout 2 - --rollout-batch-size 4 - --n-samples-per-prompt 4 - --rollout-max-response-len 8192 - --rollout-temperature 0.8 - - --global-batch-size 16 -) - -GSPO_ARGS=( - --advantage-estimator gspo - #--use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 3.5e-4 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 -) - -# launch the master node of ray in container -ray start --head --node-ip-address 127.0.0.1 --num-gpus 4 --disable-usage-stats - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}" - } - }' \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --colocate \ - --train-backend megatron \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GSPO_ARGS[@]} \ - ${VLLM_ARGS[@]} diff --git a/tests/test_hf_to_megatron.py b/tests/test_hf_to_megatron.py new file mode 100644 index 000000000..412fab08a --- /dev/null +++ b/tests/test_hf_to_megatron.py @@ -0,0 +1,573 @@ +import importlib.util +import inspect +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +# These mapping tests also run in the CPU CI image, which does not install +# Megatron. In that environment, mount the source package without executing +# megatron_utils' runtime patch initialization. +try: + _has_megatron = importlib.util.find_spec("megatron.core") is not None +except ModuleNotFoundError: + _has_megatron = False +if not _has_megatron: + _megatron_utils = types.ModuleType("vime.backends.megatron_utils") + _megatron_utils.__path__ = [str(Path(__file__).resolve().parents[1] / "vime/backends/megatron_utils")] + sys.modules["vime.backends.megatron_utils"] = _megatron_utils + +from vime.backends.megatron_utils import megatron_to_hf as megatron_to_hf_module +from vime.backends.megatron_utils.hf_to_megatron import _LOADERS +from vime.backends.megatron_utils.hf_to_megatron.common import SafetensorReader, shard_mcore_tensor +from vime.backends.megatron_utils.hf_to_megatron.deepseek import deepseek_hf_tensor +from vime.backends.megatron_utils.hf_to_megatron.glm import glm4_hf_tensor, glm4_moe_hf_tensor +from vime.backends.megatron_utils.hf_to_megatron.qwen import ( + mimo_hf_tensor, + minimax_m2_hf_tensor, + qwen_hf_tensor, + qwen_moe_hf_tensor, +) +from vime.backends.megatron_utils.hf_to_megatron.qwen3_next import qwen3_next_hf_tensor +from vime.backends.megatron_utils.hf_to_megatron.qwen3_omni import qwen3_omni_hf_tensor +from vime.backends.megatron_utils.hf_to_megatron.qwen3_vl import qwen3_vl_hf_tensor +from vime.backends.megatron_utils.megatron_to_hf import _convert_to_hf_core, convert_to_hf +from vime.backends.megatron_utils.megatron_to_hf.deepseekv3 import convert_deepseekv3_to_hf +from vime.backends.megatron_utils.megatron_to_hf.glm4 import convert_glm4_to_hf +from vime.backends.megatron_utils.megatron_to_hf.glm4moe import convert_glm4moe_to_hf +from vime.backends.megatron_utils.megatron_to_hf.mimo import convert_mimo_to_hf +from vime.backends.megatron_utils.megatron_to_hf.minimax_m2 import convert_minimax_m2_to_hf +from vime.backends.megatron_utils.megatron_to_hf.qwen2 import convert_qwen2_to_hf +from vime.backends.megatron_utils.megatron_to_hf.qwen3_next import convert_qwen3_next_to_hf +from vime.backends.megatron_utils.megatron_to_hf.qwen3_omni import convert_qwen3_omni_to_hf +from vime.backends.megatron_utils.megatron_to_hf.qwen3_vl import convert_qwen3vl_to_hf +from vime.backends.megatron_utils.megatron_to_hf.qwen3moe import convert_qwen3moe_to_hf + +NUM_GPUS = 0 + + +class Reader: + def __init__(self, **tensors): + self.tensors = tensors + + def __contains__(self, name): + return name in self.tensors + + def get_tensor(self, name): + return self.tensors[name] + + +def _config(model_type="qwen3"): + return types.SimpleNamespace( + model_type=model_type, + hidden_size=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + num_hidden_layers=2, + tie_word_embeddings=False, + ) + + +_EXPORT_ARGS = types.SimpleNamespace( + kv_channels=2, + hidden_size=8, + num_attention_heads=4, + num_query_groups=2, + num_layers=2, + q_lora_rank=None, +) + + +def test_vllm_fp8_weight_transfer_defaults_to_raw_ue8m0_scale(monkeypatch): + captured = {} + param = torch.ones(1, dtype=torch.bfloat16) + monkeypatch.setattr(megatron_to_hf_module, "remove_padding", lambda name, value, vocab_size: value) + monkeypatch.setattr( + megatron_to_hf_module, + "_convert_to_hf_core", + lambda args, model_name, name, value: [("model.weight", value)], + ) + monkeypatch.setattr( + megatron_to_hf_module, + "quantize_params", + lambda args, name, tensors, config, transform_ue8m0: captured.setdefault("transform_ue8m0", transform_ue8m0) + or tensors, + ) + + convert_to_hf(types.SimpleNamespace(vocab_size=1), "qwen3", "model.weight", param, {"quant_method": "fp8"}) + + assert captured["transform_ue8m0"] is False + assert inspect.signature(convert_to_hf).parameters["transform_ue8m0"].default is False + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("loader", "exporter", "model_type", "name", "shape"), + [ + ( + qwen_hf_tensor, + convert_qwen2_to_hf, + "qwen3", + "module.module.decoder.layers.0.self_attention.linear_qkv.weight", + (16, 8), + ), + ( + qwen_moe_hf_tensor, + convert_qwen3moe_to_hf, + "qwen3_moe", + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight3", + (12, 8), + ), + ( + mimo_hf_tensor, + convert_mimo_to_hf, + "mimo", + "module.module.mtp.layers.0.eh_proj.weight", + (8, 8), + ), + ( + minimax_m2_hf_tensor, + convert_minimax_m2_to_hf, + "minimax_m2", + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight3", + (12, 8), + ), + ( + deepseek_hf_tensor, + convert_deepseekv3_to_hf, + "deepseek_v32", + "module.module.decoder.layers.0.self_attention.wq_b.weight", + (256, 8), + ), + ( + deepseek_hf_tensor, + convert_deepseekv3_to_hf, + "deepseek_v3", + "module.module.mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight3", + (12, 8), + ), + ( + glm4_hf_tensor, + convert_glm4_to_hf, + "glm4", + "module.module.decoder.layers.0.self_attention.linear_qkv.weight", + (16, 8), + ), + ( + glm4_moe_hf_tensor, + convert_glm4moe_to_hf, + "glm4_moe", + "module.module.decoder.layers.0.mlp.shared_experts.linear_fc1.weight", + (12, 8), + ), + ( + qwen3_next_hf_tensor, + convert_qwen3_next_to_hf, + "qwen3_next", + "module.module.decoder.layers.0.self_attention.linear_qkv.weight", + (24, 8), + ), + ], +) +def test_hf_and_megatron_mappings_round_trip(loader, exporter, model_type, name, shape): + parameter = torch.arange(torch.tensor(shape).prod()).reshape(shape) + hf_tensors = dict(exporter(_EXPORT_ARGS, name, parameter)) + config = types.SimpleNamespace( + model_type=model_type, + hidden_size=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + num_hidden_layers=2, + tie_word_embeddings=False, + ) + + loaded = loader(name, Reader(**hf_tensors), config) + + assert torch.equal(loaded, parameter) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "name", + [ + "module.module.language_model.decoder.layers.0.self_attention.linear_qkv.weight", + "module.module.language_model.decoder.layers.0.mlp.experts.linear_fc1.weight3", + ], +) +def test_qwen3_omni_language_mapping_round_trip(name): + parameter = torch.arange(16 * 8).reshape(16, 8) + tensors = dict(convert_qwen3_omni_to_hf(_EXPORT_ARGS, name, parameter)) + text_config = _config("qwen3_moe") + text_config.hidden_size = 8 + text_config.num_attention_heads = 4 + text_config.num_key_value_heads = 2 + loaded = qwen3_omni_hf_tensor( + name, + Reader(**tensors), + types.SimpleNamespace(thinker_config=types.SimpleNamespace(text_config=text_config)), + ) + assert torch.equal(loaded, parameter) + + +@pytest.mark.unit +def test_qwen3_omni_encoder_mapping_is_replicated(): + parameter = torch.randn(4, 8) + name = "module.module.audio_model.layers.0.weight" + tensors = dict(convert_qwen3_omni_to_hf(_EXPORT_ARGS, name, parameter)) + loaded = qwen3_omni_hf_tensor(name, Reader(**tensors), types.SimpleNamespace()) + assert loaded is parameter + + +@pytest.mark.unit +@pytest.mark.parametrize("model_name", ["deepseekv32config", "kimik2config", "glm4moeliteconfig"]) +def test_deepseek_family_parameter_updates_use_the_direct_exporter(model_name): + parameter = torch.randn(8, 8) + + converted = _convert_to_hf_core( + _EXPORT_ARGS, + model_name, + "module.module.decoder.layers.0.self_attention.linear_proj.weight", + parameter, + ) + + assert len(converted) == 1 + assert converted[0][0] == "model.layers.0.self_attn.o_proj.weight" + assert converted[0][1] is parameter + + +@pytest.mark.unit +def test_qwen2_moe_parameter_updates_use_the_moe_exporter(): + parameter = torch.randn(12, 8) + + converted = _convert_to_hf_core( + _EXPORT_ARGS, + "qwen2_moe", + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight3", + parameter, + ) + + assert [name for name, _ in converted] == [ + "model.layers.0.mlp.experts.3.gate_proj.weight", + "model.layers.0.mlp.experts.3.up_proj.weight", + ] + assert torch.equal(converted[0][1], parameter[:6]) + assert torch.equal(converted[1][1], parameter[6:]) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "rest,shape", + [ + ("input_layernorm.weight", (8,)), + ("self_attention.linear_q_down_proj.weight", (4, 8)), + ("self_attention.linear_q_up_proj.weight", (8, 4)), + ("self_attention.linear_q_up_proj.layer_norm_weight", (4,)), + ("self_attention.linear_kv_down_proj.weight", (6, 8)), + ("self_attention.linear_kv_up_proj.weight", (8, 4)), + ("self_attention.linear_kv_up_proj.layer_norm_weight", (4,)), + ("self_attention.linear_proj.weight", (8, 8)), + ("pre_mlp_layernorm.weight", (8,)), + ("mlp.linear_fc1.weight", (12, 8)), + ("mlp.linear_fc2.weight", (8, 6)), + ("mlp.shared_experts.linear_fc1.weight", (12, 8)), + ("mlp.shared_experts.linear_fc2.weight", (8, 6)), + ("mlp.experts.linear_fc1.weight3", (12, 8)), + ("mlp.experts.linear_fc2.weight3", (8, 6)), + ("mlp.router.weight", (4, 8)), + ("mlp.router.expert_bias", (4,)), + ], +) +@pytest.mark.parametrize("mtp", [False, True]) +def test_glm_lite_native_mla_and_moe_round_trip(rest, shape, mtp): + prefix = "mtp.layers.0.transformer_layer" if mtp else "decoder.layers.1" + name = f"module.module.{prefix}.{rest}" + parameter = torch.randn(shape) + exported = _convert_to_hf_core(_EXPORT_ARGS, "glm4moeliteconfig", name, parameter) + loaded = _LOADERS["glm4_moe_lite"](name, Reader(**dict(exported)), _config("glm4_moe_lite")) + assert torch.equal(loaded, parameter) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "rest,hf_rest,shape", + [ + ("eh_proj.weight", "eh_proj.weight", (8, 16)), + ("enorm.weight", "enorm.weight", (8,)), + ("hnorm.weight", "hnorm.weight", (8,)), + ("final_layernorm.weight", "shared_head.norm.weight", (8,)), + ], +) +def test_glm_lite_native_mtp_projection_and_norm_round_trip(rest, hf_rest, shape): + name = f"module.module.mtp.layers.0.{rest}" + parameter = torch.randn(shape) + exported = _convert_to_hf_core(_EXPORT_ARGS, "glm4moeliteconfig", name, parameter) + assert [key for key, _ in exported] == [f"model.layers.{_EXPORT_ARGS.num_layers}.{hf_rest}"] + loaded = _LOADERS["glm4_moe_lite"](name, Reader(**dict(exported)), _config("glm4_moe_lite")) + assert torch.equal(loaded, parameter) + + +@pytest.mark.unit +def test_qwen_and_llama_share_the_basic_qkv_mapping(): + q = torch.arange(16).view(4, 4) + k = torch.arange(8).view(2, 4) + 100 + v = torch.arange(8).view(2, 4) + 200 + reader = Reader( + **{ + "model.layers.3.self_attn.q_proj.weight": q, + "model.layers.3.self_attn.k_proj.weight": k, + "model.layers.3.self_attn.v_proj.weight": v, + } + ) + + loaded = qwen_hf_tensor( + "module.module.decoder.layers.3.self_attention.linear_qkv.weight", + reader, + _config(), + ) + + assert torch.equal(loaded, torch.cat((q, k, v))) + assert _LOADERS["qwen3"] is _LOADERS["llama"] is qwen_hf_tensor + + +@pytest.mark.unit +def test_qwen_moe_merges_one_global_expert(): + gate = torch.randn(4, 3) + up = torch.randn(4, 3) + reader = Reader( + **{ + "model.layers.1.mlp.experts.9.gate_proj.weight": gate, + "model.layers.1.mlp.experts.9.up_proj.weight": up, + } + ) + + loaded = qwen_moe_hf_tensor( + "module.module.decoder.layers.1.mlp.experts.linear_fc1.weight9", + reader, + _config(), + ) + + assert torch.equal(loaded, torch.cat((gate, up))) + + +@pytest.mark.unit +def test_deepseek_mapping_handles_kimi_and_dsa_layouts(): + mla = torch.randn(8, 4) + kimi = deepseek_hf_tensor( + "module.module.decoder.layers.0.self_attention.linear_kv_down_proj.weight", + Reader(**{"model.layers.0.self_attn.kv_a_proj_with_mqa.weight": mla}), + _config("kimi_k2"), + ) + assert kimi is mla + + dsa = torch.arange(128 * 2).view(128, 2) + reordered = deepseek_hf_tensor( + "module.module.decoder.layers.0.self_attention.wk.weight", + Reader(**{"model.layers.0.self_attn.indexer.wk.weight": dsa}), + _config("glm_moe_dsa"), + ) + assert torch.equal(reordered, torch.cat((dsa[64:], dsa[:64]))) + + +@pytest.mark.unit +def test_glm_dense_and_moe_mtp_use_native_mappings(): + fused = torch.randn(8, 4) + dense = glm4_hf_tensor( + "module.module.decoder.layers.1.mlp.linear_fc1.weight", + Reader(**{"model.layers.1.mlp.gate_up_proj.weight": fused}), + _config("glm4"), + ) + assert dense is fused + + mtp = torch.randn(4, 4) + moe = glm4_moe_hf_tensor( + "module.module.mtp.layers.0.eh_proj.weight", + Reader(**{"model.layers.2.eh_proj.weight": mtp}), + _config("glm4_moe"), + ) + assert moe is mtp + + gate = torch.randn(4, 3) + up = torch.randn(4, 3) + shared = glm4_moe_hf_tensor( + "module.module.decoder.layers.0.mlp.shared_experts.linear_fc1.weight", + Reader( + **{ + "model.layers.0.mlp.shared_experts.gate_proj.weight": gate, + "model.layers.0.mlp.shared_experts.up_proj.weight": up, + } + ), + _config("glm4_moe"), + ) + assert torch.equal(shared, torch.cat((gate, up))) + + +@pytest.mark.unit +def test_minimax_and_mimo_keep_their_small_qwen_deltas(): + gate = torch.randn(4, 3) + up = torch.randn(4, 3) + minimax = minimax_m2_hf_tensor( + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight2", + Reader( + **{ + "model.layers.0.block_sparse_moe.experts.2.w1.weight": gate, + "model.layers.0.block_sparse_moe.experts.2.w3.weight": up, + } + ), + _config("minimax_m2"), + ) + assert torch.equal(minimax, torch.cat((gate, up))) + + hf_eh = torch.arange(24).view(3, 8) + mimo = mimo_hf_tensor( + "module.module.mtp.layers.0.eh_proj.weight", + Reader(**{"model.mtp_layers.0.input_proj.weight": hf_eh}), + _config("mimo"), + ) + assert torch.equal(mimo, torch.cat((hf_eh[:, 4:], hf_eh[:, :4]), dim=1)) + + +@pytest.mark.unit +def test_loader_scope_stays_explicit(): + assert set(_LOADERS) == { + "deepseek_v3", + "deepseek_v32", + "glm4", + "glm4_moe", + "glm4_moe_lite", + "glm_moe_dsa", + "kimi_k2", + "llama", + "mimo", + "minimax_m2", + "qwen2", + "qwen2_moe", + "qwen3", + "qwen3_5", + "qwen3_5_moe", + "qwen3_moe", + "qwen3_next", + "qwen3_omni_moe", + "qwen3_vl", + } + + +@pytest.mark.unit +@pytest.mark.parametrize(("model_name", "prefix"), [("qwen3", ""), ("qwen3_vl", "language_model.")]) +@pytest.mark.parametrize( + ("name", "shape", "is_vocab"), + [ + ("embedding.word_embeddings.weight", (16, 8), True), + ("output_layer.weight", (16, 8), True), + ("decoder.final_layernorm.weight", (16,), False), + ], +) +def test_native_export_removes_only_vocab_padding(model_name, prefix, name, shape, is_vocab): + args = types.SimpleNamespace(vocab_size=8) + weight = torch.randn(shape) + [(_, exported)] = convert_to_hf(args, model_name, "module.module." + prefix + name, weight) + expected = weight[: args.vocab_size] if is_vocab else weight + assert torch.equal(exported, expected) + if not is_vocab: + assert exported is weight + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("name", "shape"), + [ + ("embedding.word_embeddings.weight", (16, 8)), + ("output_layer.weight", (16, 8)), + ("decoder.final_layernorm.weight", (8,)), + ("decoder.layers.0.self_attention.linear_qkv.weight", (16, 8)), + ("decoder.layers.0.self_attention.linear_qkv.bias", (16,)), + ("decoder.layers.0.self_attention.linear_proj.weight", (8, 8)), + ("decoder.layers.0.self_attention.linear_qkv.layer_norm_weight", (8,)), + ("decoder.layers.0.self_attention.q_layernorm.weight", (2,)), + ("decoder.layers.0.self_attention.k_layernorm.weight", (2,)), + ("decoder.layers.0.mlp.linear_fc1.weight", (24, 8)), + ("decoder.layers.0.mlp.linear_fc2.weight", (8, 12)), + ("decoder.layers.0.mlp.linear_fc1.layer_norm_weight", (8,)), + ], +) +def test_qwen3_vl_language_round_trip(name, shape): + name = "module.module.language_model." + name + weight = torch.randn(shape) + config = types.SimpleNamespace( + tie_word_embeddings=False, + text_config=types.SimpleNamespace( + hidden_size=8, num_attention_heads=4, num_key_value_heads=2, head_dim=2, tie_word_embeddings=False + ), + ) + reader = Reader(**dict(convert_qwen3vl_to_hf(_EXPORT_ARGS, name, weight))) + assert torch.equal(qwen3_vl_hf_tensor(name, reader, config), weight) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("top_tied", "text_tied", "has_head"), [(True, False, True), (False, True, True), (False, False, False)] +) +def test_qwen3_vl_tied_output_uses_language_embedding(top_tied, text_tied, has_head): + embedding = torch.randn(8, 4) + tensors = {"model.language_model.embed_tokens.weight": embedding} + if has_head: + tensors["lm_head.weight"] = torch.zeros_like(embedding) + config = types.SimpleNamespace( + tie_word_embeddings=top_tied, text_config=types.SimpleNamespace(tie_word_embeddings=text_tied) + ) + assert qwen3_vl_hf_tensor("output_layer.weight", Reader(**tensors), config) is embedding + + +@pytest.mark.unit +@pytest.mark.parametrize("platform", ["cuda", "npu"]) +@pytest.mark.parametrize("expert", [False, True]) +@pytest.mark.parametrize("parallel_size", [1, 2, 4]) +def test_native_fc1_loading_uses_platform_partition_metadata(monkeypatch, platform, expert, parallel_size): + mpu = pytest.importorskip("megatron.core").mpu + monkeypatch.setenv("VIME_PLATFORM", platform) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: 8 if expert else parallel_size) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0 if expert else parallel_size - 1) + monkeypatch.setattr(mpu, "get_expert_tensor_parallel_world_size", lambda: parallel_size if expert else 8) + monkeypatch.setattr(mpu, "get_expert_tensor_parallel_rank", lambda: parallel_size - 1 if expert else 0) + name = "decoder.layers.0.mlp." + ("experts.linear_fc1.weight0" if expert else "linear_fc1.weight") + weight = torch.arange(16 * 8).reshape(16, 8) + # MindSpeed grouped column-parallel weights carry dim=1, but store [out, in]. + parameter = types.SimpleNamespace( + tensor_model_parallel=True, partition_dim=1 if platform == "npu" and expert else 0, partition_stride=1 + ) + shard = shard_mcore_tensor(name, weight, parameter) + gate, up = weight.chunk(2) + assert torch.equal(shard, torch.cat((gate.chunk(parallel_size)[-1], up.chunk(parallel_size)[-1]))) + + +@pytest.mark.unit +@pytest.mark.parametrize("parallel_mode", [None, "duplicated"]) +def test_native_loading_does_not_shard_replicated_parameters(monkeypatch, parallel_mode): + pytest.importorskip("megatron.core") + monkeypatch.setenv("VIME_PLATFORM", "npu") + parameter = types.SimpleNamespace(tensor_model_parallel=parallel_mode is not None, parallel_mode=parallel_mode) + weight = torch.randn(4, 8) + assert shard_mcore_tensor("model.visual.blocks.0.mlp.linear_fc1.weight", weight, parameter) is weight + + +@pytest.mark.unit +def test_reader_dequantizes_block_scaled_fp8(tmp_path): + weight = torch.linspace(-2, 2, 128 * 128).view(128, 128).to(torch.float8_e4m3fn) + scale = torch.tensor([[2.0]]) + save_file( + {"weight": weight, "weight_scale_inv": scale}, + tmp_path / "model.safetensors", + ) + + loaded = SafetensorReader(tmp_path).get_tensor("weight") + + assert loaded.dtype == torch.bfloat16 + assert torch.equal(loaded, weight.to(torch.bfloat16) * 2) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_layerwise_alignment.py b/tests/test_layerwise_alignment.py new file mode 100644 index 000000000..ec452577e --- /dev/null +++ b/tests/test_layerwise_alignment.py @@ -0,0 +1,79 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vime.backends.megatron_utils.alignment.layerwise_alignment import enable_megatron_layerwise_dump + +NUM_GPUS = 0 + + +class _Layer(nn.Module): + def __init__(self, layer_number: int): + super().__init__() + self.layer_number = layer_number + + def forward(self, value): + return value + self.layer_number + + +class _Decoder(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Layer(1), _Layer(2)]) + + def forward(self, value): + for layer in self.layers: + value = layer(value) + return value + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.decoder = _Decoder() + + def forward(self, *, input_ids, packed_seq_params): + del packed_seq_params + return self.decoder(input_ids.float()) + + +def test_megatron_layerwise_dump(monkeypatch, tmp_path): + monkeypatch.setenv("VIME_LAYERWISE_ALIGNMENT_DUMP_DIR", str(tmp_path)) + args = SimpleNamespace(megatron_deepgemm_forward_layers=[0, 1]) + model = _Model() + enable_megatron_layerwise_dump(args, [model], store_prefix="actor_") + + model( + input_ids=torch.tensor([[7, 8]]), + packed_seq_params=SimpleNamespace(cu_seqlens_q=torch.tensor([0, 2])), + ) + + (dump_file,) = list(tmp_path.glob("rank*/actor_Pass*.pt")) + values = torch.load(dump_file, weights_only=False) + torch.testing.assert_close(values["input_ids"], torch.tensor([[7, 8]])) + torch.testing.assert_close(values["layers"][0], torch.tensor([[8.0, 9.0]])) + torch.testing.assert_close(values["layers"][1], torch.tensor([[10.0, 11.0]])) + + +def test_megatron_layerwise_dump_is_enabled_on_nonzero_rank(monkeypatch, tmp_path): + monkeypatch.setenv("VIME_LAYERWISE_ALIGNMENT_DUMP_DIR", str(tmp_path)) + monkeypatch.setenv("VIME_LAYERWISE_ALIGNMENT_MODULE_SUFFIXES", "decoder.layers.0") + monkeypatch.setattr("vime.backends.megatron_utils.alignment.layerwise_alignment._global_rank", lambda: 3) + args = SimpleNamespace(megatron_deepgemm_forward_layers=[0, 1]) + model = _Model() + + enable_megatron_layerwise_dump(args, [model], store_prefix="actor_") + model( + input_ids=torch.tensor([[7, 8]]), + packed_seq_params=SimpleNamespace(cu_seqlens_q=torch.tensor([0, 2])), + ) + + (dump_file,) = list(tmp_path.glob("rank00003/actor_Pass*.pt")) + values = torch.load(dump_file, weights_only=False) + assert "decoder.layers.0" in values["modules"] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_logprob_response_spans.py b/tests/test_logprob_response_spans.py new file mode 100644 index 000000000..d22ef4591 --- /dev/null +++ b/tests/test_logprob_response_spans.py @@ -0,0 +1,99 @@ +from argparse import Namespace + +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from megatron.core import mpu +from vime.backends.megatron_utils.loss import _build_topp_keep_mask, get_rollout_top_p_logprob_kwargs + + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_missing_top_p_replay_data_raises(): + with pytest.raises(ValueError, match="requires rollout_top_p_token_ids"): + get_rollout_top_p_logprob_kwargs(Namespace(rollout_top_p=0.95), {}) + + +def _set_cp(monkeypatch, *, size: int, rank: int) -> None: + monkeypatch.setattr(mpu, "get_context_parallel_world_size", lambda: size) + monkeypatch.setattr(mpu, "get_context_parallel_rank", lambda: rank) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0, raising=False) + + +def _kept_ids(row: torch.Tensor) -> list[int]: + return row.nonzero(as_tuple=False).squeeze(-1).tolist() + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("rank", "expected"), + [ + (0, {2: [107]}), + (1, {1: [104], 2: [105], 3: [106]}), + ], +) +def test_top_p_mask_aligns_with_zigzag_cp_response_rows(monkeypatch, rank, expected): + _set_cp(monkeypatch, size=2, rank=rank) + keep = _build_topp_keep_mask( + 4, + 200, + torch.device("cpu"), + top_p_token_ids=[[104, 105, 106, 107]], + top_p_token_offsets=[[0, 1, 2, 3, 4]], + total_lengths=[8], + response_lengths=[4], + allgather_cp=False, + ) + + masked_rows = {row: _kept_ids(keep[row]) for row in range(keep.size(0)) if not keep[row].all()} + assert masked_rows == expected + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("rank", "expected"), + [ + (0, {1: [102], 2: [103]}), + (1, {0: [104], 1: [105]}), + ], +) +def test_top_p_mask_aligns_with_allgather_cp_response_rows(monkeypatch, rank, expected): + _set_cp(monkeypatch, size=2, rank=rank) + keep = _build_topp_keep_mask( + 3, + 200, + torch.device("cpu"), + top_p_token_ids=[[102, 103, 104, 105]], + top_p_token_offsets=[[0, 1, 2, 3, 4]], + total_lengths=[6], + response_lengths=[4], + allgather_cp=True, + ) + + masked_rows = {row: _kept_ids(keep[row]) for row in range(keep.size(0)) if not keep[row].all()} + assert masked_rows == expected + + +@pytest.mark.unit +def test_top_p_mask_aligns_with_cp1_response_rows(monkeypatch): + _set_cp(monkeypatch, size=1, rank=0) + keep = _build_topp_keep_mask( + 9, + 30, + torch.device("cpu"), + top_p_token_ids=[[13, 99, 14], [21, 22, 99, 23]], + top_p_token_offsets=[[0, 2, 3], [0, 1, 3, 4]], + total_lengths=[5, 4], + response_lengths=[2, 3], + allgather_cp=False, + ) + + masked_rows = {row: _kept_ids(keep[row]) for row in range(keep.size(0)) if not keep[row].all()} + assert masked_rows == {2: [13], 3: [14], 5: [21], 6: [22], 7: [23]} + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_loss_cp_invariance.py b/tests/test_loss_cp_invariance.py index f28679aac..6025de75c 100644 --- a/tests/test_loss_cp_invariance.py +++ b/tests/test_loss_cp_invariance.py @@ -48,7 +48,7 @@ The contract here is on *our* scaling math (steps 1 + 4 are vime's; step 2 is what Megatron does to our 3-tuple). If Megatron later changes step 2 — e.g. drops the ``/= num_microbatches`` — this test won't catch -it, but the real GPU integration suite (``test_qwen2.5_0.5B_short.py``) +it, but the real GPU parallelism suite (``test_qwen3_0.6B_parallel_check.py``) will. """ diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index af20fa71c..fd1247191 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -1,3 +1,4 @@ +import argparse import importlib.util import sys import types @@ -5,7 +6,6 @@ import pytest - NUM_GPUS = 0 @@ -39,6 +39,35 @@ def load_arguments_module(monkeypatch): return module +def load_vime_arguments_module(monkeypatch): + router_pkg_mod = types.ModuleType("vllm_router") + router_launch_mod = types.ModuleType("vllm_router.launch_router") + vllm_arguments_mod = types.ModuleType("vime.backends.vllm_utils.arguments") + vllm_external_mod = types.ModuleType("vime.backends.vllm_utils.external") + logging_utils_mod = types.ModuleType("vime.observability.logging_utils") + + router_launch_mod.RouterArgs = object + vllm_arguments_mod.vllm_parse_args = lambda *args, **kwargs: None + vllm_arguments_mod.validate_args = lambda args: args + vllm_external_mod.apply_external_engine_info_to_args = lambda *args, **kwargs: None + logging_utils_mod.configure_logger = lambda *args, **kwargs: None + + monkeypatch.setitem(sys.modules, "vllm_router", router_pkg_mod) + monkeypatch.setitem(sys.modules, "vllm_router.launch_router", router_launch_mod) + monkeypatch.setitem(sys.modules, "vime.backends.vllm_utils.arguments", vllm_arguments_mod) + monkeypatch.setitem(sys.modules, "vime.backends.vllm_utils.external", vllm_external_mod) + monkeypatch.setitem(sys.modules, "vime.observability.logging_utils", logging_utils_mod) + + module_path = Path(__file__).resolve().parents[1] / "vime" / "utils" / "arguments.py" + module_name = "test_vime_argument_validation_module" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + def make_qwen3_6_args(**overrides): values = dict( hidden_size=2048, @@ -139,5 +168,267 @@ def test_allgather_cp_ignores_cp_size_one(monkeypatch): module._validate_allgather_cp_supported(args) +@pytest.mark.unit +def test_update_weight_disk_dir_required_for_disk_transport(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(update_weight_transport="disk", update_weight_disk_dir=None) + + with pytest.raises(ValueError, match="update-weight-disk-dir"): + module.vime_validate_args(args) + + +def make_vime_validate_args(**overrides): + values = dict( + eval_config=None, + eval_prompt_data=None, + kl_coef=0, + use_kl_loss=False, + ref_load=None, + use_opd=False, + opd_type=None, + opd_teacher_load=None, + load=None, + hf_checkpoint="/tmp/hf", + ref_ckpt_step=None, + ckpt_step=None, + no_load_optim=False, + no_load_rng=False, + finetune=False, + start_rollout_id=None, + eval_interval=None, + save_interval=None, + save=None, + kl_loss_coef=0, + advantage_estimator="grpo", + normalize_advantages=False, + use_rollout_logprobs=False, + use_tis=False, + get_mismatch_metrics=False, + custom_tis_function_path=None, + use_dynamic_batch_size=False, + max_tokens_per_gpu=None, + log_probs_max_tokens_per_gpu=None, + balance_by_flops=False, + balance_data=False, + eps_clip_high=None, + eps_clip=0.2, + eval_reward_key=None, + reward_key="reward", + dump_details=None, + save_debug_rollout_data=None, + save_debug_train_data=None, + load_debug_rollout_data=None, + rollout_external_engine_addrs=None, + debug_train_only=False, + actor_num_gpus_per_node=8, + actor_num_nodes=1, + num_gpus_per_node=8, + offload=False, + offload_train=None, + offload_rollout=None, + debug_rollout_only=False, + colocate=False, + rollout_num_gpus=8, + eval_function_path=None, + rollout_function_path="custom.rollout", + vllm_speculative_config=None, + num_steps_per_rollout=None, + rollout_batch_size=1, + n_samples_per_prompt=1, + global_batch_size=None, + grpo_std_normalization=True, + over_sampling_batch_size=None, + num_epoch=None, + num_rollout=1, + rollout_global_dataset=False, + enable_mtp_training=False, + mtp_num_layers=None, + use_rollout_routing_replay=False, + use_routing_replay=False, + custom_config_path=None, + eval_max_context_len=None, + rollout_max_context_len=None, + rollout_max_prompt_len=None, + train_backend="megatron", + release_train=False, + keep_old_actor=False, + only_train_params_name_list=None, + freeze_params_name_list=None, + update_weight_transport="nccl", + update_weight_disk_dir=None, + update_weight_local_checkpoint_dir=None, + update_weight_mode="full", + rollout_temperature=1.0, + ) + values.update(overrides) + return types.SimpleNamespace(**values) + + +@pytest.mark.unit +def test_vime_validate_args_preserves_explicit_start_rollout_id(monkeypatch): + """``--start-rollout-id`` is only a fallback when the user did not set it. + + An explicit value is needed when there is no resumable Megatron checkpoint. + """ + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(start_rollout_id=100) + + module.vime_validate_args(args) + + assert args.start_rollout_id == 100 + + +@pytest.mark.unit +def test_vime_validate_args_defaults_start_rollout_id_to_zero(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(start_rollout_id=None) + + module.vime_validate_args(args) + + assert args.start_rollout_id == 0 + + +@pytest.mark.unit +def test_vime_validate_args_derives_dspark_from_speculative_method(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(vllm_speculative_config={"method": "dspark"}) + + module.vime_validate_args(args) + + assert args.dspark_enabled is True + + +@pytest.mark.unit +def test_vime_validate_args_rejects_equal_debug_data_paths(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + save_debug_rollout_data="/tmp/debug_{rollout_id}.pt", + save_debug_train_data="/tmp/debug_{rollout_id}.pt", + ) + + with pytest.raises(ValueError, match="--save-debug-train-data must not be equal"): + module.vime_validate_args(args) + + +@pytest.mark.unit +@pytest.mark.parametrize("temperature", [0.0, -0.1]) +def test_vime_validate_args_rejects_non_positive_rollout_temperature(monkeypatch, temperature): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(rollout_temperature=temperature) + + with pytest.raises(ValueError, match="--rollout-temperature must be > 0"): + module.vime_validate_args(args) + + +@pytest.mark.unit +def test_vime_validate_args_preserves_zero_rollout_gpus_under_colocate(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(colocate=True, rollout_num_gpus=0) + + module.vime_validate_args(args) + + assert args.rollout_num_gpus == 0 + assert args.offload_train is True + assert args.offload_rollout is True + + +@pytest.mark.unit +def test_vime_validate_args_preserves_larger_rollout_gpus_under_colocate(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + colocate=True, + actor_num_gpus_per_node=8, + actor_num_nodes=1, + rollout_num_gpus=12, + ) + + module.vime_validate_args(args) + + assert args.rollout_num_gpus == 12 + assert args.offload_train is True + assert args.offload_rollout is True + + +@pytest.mark.unit +def test_vime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args(colocate=False, rollout_num_gpus=0) + + module.vime_validate_args(args) + + assert args.rollout_num_gpus == 0 + assert args.actor_num_gpus_per_node == 8 + assert args.actor_num_nodes == 1 + assert args.offload_train is False + assert args.offload_rollout is False + + +@pytest.mark.unit +def test_update_weight_delta_disk_is_valid(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + update_weight_mode="delta", + update_weight_transport="disk", + update_weight_disk_dir="/shared/delta", + update_weight_local_checkpoint_dir="/local/delta", + ) + + module.vime_validate_args(args) + + +@pytest.mark.unit +def test_update_weight_delta_requires_disk_transport(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + update_weight_mode="delta", + update_weight_transport="nccl", + update_weight_local_checkpoint_dir="/local/delta", + ) + + with pytest.raises(ValueError, match="requires --update-weight-transport=disk"): + module.vime_validate_args(args) + + +@pytest.mark.unit +def test_update_weight_delta_rejects_colocate(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + update_weight_mode="delta", + update_weight_transport="disk", + update_weight_disk_dir="/shared/delta", + update_weight_local_checkpoint_dir="/local/delta", + colocate=True, + ) + + with pytest.raises(ValueError, match="not supported with --colocate"): + module.vime_validate_args(args) + + +@pytest.mark.unit +def test_update_weight_delta_requires_local_checkpoint_dir(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = make_vime_validate_args( + update_weight_mode="delta", + update_weight_transport="disk", + update_weight_disk_dir="/shared/delta", + ) + + with pytest.raises(ValueError, match="requires --update-weight-local-checkpoint-dir"): + module.vime_validate_args(args) + + +@pytest.mark.unit +def test_force_fp8_ue8m0_scale_argument(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + parser = argparse.ArgumentParser() + module.get_vime_extra_args_provider()(parser) + + defaults = parser.parse_args(["--rollout-batch-size", "1"]) + configured = parser.parse_args(["--rollout-batch-size", "1", "--force-fp8-ue8m0-scale"]) + + assert defaults.force_fp8_ue8m0_scale is False + assert configured.force_fp8_ue8m0_scale is True + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_metric_report.py b/tests/test_metric_report.py index 9ab724f45..f0c67713a 100644 --- a/tests/test_metric_report.py +++ b/tests/test_metric_report.py @@ -1,8 +1,7 @@ """Single-process metric-report invariance tests. Pins train-side / rollout-side report formulas implemented in -``vime.backends.megatron_utils.cp_utils.reduce_train_step_metrics`` and -``rollout_log_metric_contribution``: the reported number for a given set +``vime.observability.train_metric_utils``: the reported number for a given set of samples must be the same regardless of - how samples are distributed across micro-batches / DP ranks @@ -28,11 +27,12 @@ from vime.backends.megatron_utils.cp_utils import ( # noqa: E402 get_logits_and_tokens_offset_with_cp, get_sum_of_sample_mean, +) +from vime.observability.train_metric_utils import ( # noqa: E402 reduce_train_step_metrics, rollout_log_metric_contribution, ) - NUM_GPUS = 0 @@ -143,7 +143,7 @@ def _simulate_rollout_report(samples_per_rank): per-token metric branch. Each "rank" applies the reducer once over its full sample subset, then - ``rollout_log_metric_contribution`` (the same helper data.py uses) emits + ``rollout_log_metric_contribution`` (the same helper the reporter uses) emits the ``(per_rank_sum, count)`` tuple. We aggregate via ``Σsum / Σcount`` — the same shape ``gather_log_data`` uses. """ diff --git a/tests/test_metric_report_dist.py b/tests/test_metric_report_dist.py index aca4e3fae..132e483ad 100644 --- a/tests/test_metric_report_dist.py +++ b/tests/test_metric_report_dist.py @@ -1,4 +1,4 @@ -"""Multi-process distributed tests for the cp_utils report helpers. +"""Multi-process distributed tests for the train metric report helpers. Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` (gloo backend) and exercise the actual production helpers end-to-end. The @@ -42,7 +42,6 @@ stub_megatron_in_worker, ) - NUM_GPUS = 0 @@ -68,7 +67,8 @@ def _train_step_distributed_worker( # Import AFTER the megatron stub override so cp_utils still binds # against the pre-installed stub (which we've now pinned for this # worker's CP rank). - from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean, reduce_train_step_metrics + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean + from vime.observability.train_metric_utils import reduce_train_step_metrics all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS @@ -197,11 +197,8 @@ def _rollout_log_distributed_worker( dp_group = init_worker_process_group(rank, world_size, master_port) try: - from vime.backends.megatron_utils.cp_utils import ( - gather_and_reduce_log_dict, - get_sum_of_sample_mean, - rollout_log_metric_contribution, - ) + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean + from vime.observability.train_metric_utils import gather_and_reduce_log_dict, rollout_log_metric_contribution all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS diff --git a/tests/test_mimo_7B_mtp_only_grad.py b/tests/test_mimo_7B_mtp_only_grad.py index ec315f043..2c49cd911 100644 --- a/tests/test_mimo_7B_mtp_only_grad.py +++ b/tests/test_mimo_7B_mtp_only_grad.py @@ -91,8 +91,9 @@ def execute(): "--rollout-num-gpus-per-engine 2 " "--rollout-num-gpus 8 " "--vllm-gpu-memory-utilization 0.8 " - "--vllm-max-cudagraph-capture-size 8 " - '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":2}\' ' + "--vllm-enforce-eager " + "--vllm-max-cudagraph-capture-size 32 " + '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":3}\' ' ) # Enable MTP training with loss scaling diff --git a/tests/test_model_provider_freeze.py b/tests/test_model_provider_freeze.py new file mode 100644 index 000000000..e1bf45f0b --- /dev/null +++ b/tests/test_model_provider_freeze.py @@ -0,0 +1,127 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + +NUM_GPUS = 0 + + +def _load_model_provider(monkeypatch): + modules = { + "megatron": types.ModuleType("megatron"), + "megatron.core": types.ModuleType("megatron.core"), + "megatron.core.models": types.ModuleType("megatron.core.models"), + "megatron.core.models.gpt": types.ModuleType("megatron.core.models.gpt"), + "megatron.core.models.gpt.gpt_layer_specs": types.ModuleType("megatron.core.models.gpt.gpt_layer_specs"), + "megatron.core.transformer": types.ModuleType("megatron.core.transformer"), + "megatron.core.transformer.spec_utils": types.ModuleType("megatron.core.transformer.spec_utils"), + "megatron.core.transformer.transformer_config": types.ModuleType( + "megatron.core.transformer.transformer_config" + ), + "megatron.training": types.ModuleType("megatron.training"), + "megatron.training.arguments": types.ModuleType("megatron.training.arguments"), + "vime.utils.misc": types.ModuleType("vime.utils.misc"), + } + modules["megatron.core"].tensor_parallel = types.SimpleNamespace() + modules["megatron.core.models.gpt"].GPTModel = torch.nn.Module + layer_specs = modules["megatron.core.models.gpt.gpt_layer_specs"] + layer_specs.get_gpt_decoder_block_spec = lambda *args, **kwargs: None + layer_specs.get_gpt_layer_local_spec = lambda *args, **kwargs: None + layer_specs.get_gpt_layer_with_transformer_engine_spec = lambda *args, **kwargs: None + modules["megatron.core.transformer.spec_utils"].import_module = lambda value: value + modules["megatron.core.transformer.transformer_config"].TransformerConfig = object + modules["megatron.training.arguments"].core_transformer_config_from_args = lambda args: object() + modules["vime.utils.misc"].load_function = lambda value: value + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + module_path = Path(__file__).resolve().parents[1] / "vime" / "backends" / "megatron_utils" / "model_provider.py" + module_name = "test_model_provider_freeze_module" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class _GLMIndexerAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.wq_b = torch.nn.Linear(2, 2, bias=False) + self.wk = torch.nn.Linear(2, 2, bias=False) + self.k_norm = torch.nn.LayerNorm(2) + self.weights_proj = torch.nn.Linear(2, 2, bias=False) + self.linear_q_down_proj = torch.nn.Linear(2, 2, bias=False) + + +class _UpstreamDSAAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.core_attention = torch.nn.Module() + self.core_attention.indexer = torch.nn.Module() + self.core_attention.indexer.linear_wq_b = torch.nn.Linear(2, 2, bias=False) + self.core_attention.indexer.linear_wk = torch.nn.Linear(2, 2, bias=False) + self.core_attention.indexer.k_norm = torch.nn.LayerNorm(2) + self.core_attention.indexer.linear_weights_proj = torch.nn.Linear(2, 2, bias=False) + self.core_attention.regular_projection = torch.nn.Linear(2, 2, bias=False) + + +class _Layer(torch.nn.Module): + def __init__(self, attention): + super().__init__() + self.self_attention = attention + self.mlp = torch.nn.Linear(2, 2, bias=False) + + +class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_Layer(_GLMIndexerAttention()), _Layer(_UpstreamDSAAttention())]) + self.output_layer = torch.nn.Linear(2, 2, bias=False) + + +@pytest.mark.unit +def test_freeze_indexer_covers_glm_and_upstream_dsa_names(monkeypatch): + model_provider = _load_model_provider(monkeypatch) + model = _Model() + args = types.SimpleNamespace( + only_train_params_name_list=None, + freeze_params_name_list=None, + freeze_indexer=True, + ) + + model_provider.freeze_model_params(model, args) + + frozen = set(model._vime_frozen_indexer_param_names) + assert frozen + for name, parameter in model.named_parameters(): + if name in frozen: + assert not parameter.requires_grad, name + else: + assert parameter.requires_grad, name + assert "layers.0.self_attention.linear_q_down_proj.weight" not in frozen + assert "layers.1.self_attention.core_attention.regular_projection.weight" not in frozen + assert "layers.0.mlp.weight" not in frozen + assert "output_layer.weight" not in frozen + + +@pytest.mark.unit +def test_freeze_indexer_rejects_unrecognized_attention(monkeypatch): + model_provider = _load_model_provider(monkeypatch) + model = _Layer(torch.nn.MultiheadAttention(2, 1)) + args = types.SimpleNamespace( + only_train_params_name_list=None, + freeze_params_name_list=None, + freeze_indexer=True, + ) + + with pytest.raises(RuntimeError, match="no recognized DSA indexer"): + model_provider.freeze_model_params(model, args) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_moonlight_16B_A3B_non_colocate_nccl.py b/tests/test_moonlight_16B_A3B_non_colocate_nccl.py new file mode 100644 index 000000000..3bfd180df --- /dev/null +++ b/tests/test_moonlight_16B_A3B_non_colocate_nccl.py @@ -0,0 +1,127 @@ +"""Eight-GPU non-colocated MoE E2E for native vLLM NCCL updates.""" + +import os + +import vime.utils.external_utils.command_utils as U + + +os.environ.setdefault("NCCL_NVLS_ENABLE", "0") + + +MODEL_NAME = "Moonlight-16B-A3B-Instruct" +MODEL_TYPE = "moonlight" +NUM_GPUS = 8 +TRAIN_GPUS = 4 +ROLLOUT_GPUS = 4 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command( + "hf download moonshotai/Moonlight-16B-A3B-Instruct " "--local-dir /root/models/Moonlight-16B-A3B-Instruct" + ) + U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=TRAIN_GPUS, + dir_dst="/dev/shm", + ) + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} " f"--ref-load {TORCH_DIST_CKPT} " + + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 2 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 512 " + "--rollout-temperature 1.0 " + "--global-batch-size 8 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + vllm_args = ( + f"--rollout-num-gpus {ROLLOUT_GPUS} " + f"--rollout-num-gpus-per-engine {ROLLOUT_GPUS} " + f"--vllm-data-parallel-size {ROLLOUT_GPUS} " + "--vllm-enable-expert-parallel " + "--vllm-all2all-backend naive " + "--vllm-gpu-memory-utilization 0.65 " + "--vllm-max-model-len 4096 " + "--vllm-max-num-seqs 8 " + "--vllm-enforce-eager " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {TRAIN_GPUS} " + "--ci-test " + ) + + U.execute_train( + train_args=( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{perf_args} " + f"{vllm_args} " + f"{misc_args} " + ), + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_placement_group.py b/tests/test_placement_group.py new file mode 100644 index 000000000..55550e21d --- /dev/null +++ b/tests/test_placement_group.py @@ -0,0 +1,54 @@ +import sys +from argparse import Namespace +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.ray.placement_group import _create_placement_group, _get_placement_group_layout + +NUM_GPUS = 0 + + +def _args(**overrides): + values = { + "actor_num_nodes": 2, + "actor_num_gpus_per_node": 8, + "rollout_num_gpus": 32, + "debug_train_only": False, + "debug_rollout_only": False, + "colocate": False, + "rollout_external": False, + } + values.update(overrides) + return Namespace(**values) + + +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + pytest.param({}, (48, 16), id="normal_non_colocate"), + pytest.param({"debug_train_only": True}, (16, 0), id="debug_train_only"), + pytest.param({"debug_rollout_only": True}, (32, 0), id="debug_rollout_only"), + pytest.param({"colocate": True, "rollout_num_gpus": 8}, (16, 0), id="colocate_rollout_less_than_actor"), + pytest.param({"colocate": True, "rollout_num_gpus": 16}, (16, 0), id="colocate_rollout_equals_actor"), + pytest.param({"colocate": True, "rollout_num_gpus": 32}, (32, 0), id="colocate_rollout_more_than_actor"), + pytest.param({"rollout_num_gpus": 0}, (16, 16), id="zero_rollout_gpus"), + pytest.param({"colocate": True, "rollout_num_gpus": 0}, (16, 0), id="colocate_zero_rollout_gpus"), + pytest.param({"rollout_external": True}, (16, 16), id="external"), + pytest.param({"rollout_external": True, "debug_rollout_only": True}, (16, 0), id="external_debug_rollout"), + ], +) +def test_placement_group_layout(overrides, expected): + assert _get_placement_group_layout(_args(**overrides)) == expected + + +def test_create_zero_gpu_placement_group_is_empty(): + assert _create_placement_group(0) == (None, [], []) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_policy_loss.py b/tests/test_policy_loss.py new file mode 100644 index 000000000..1e45925b5 --- /dev/null +++ b/tests/test_policy_loss.py @@ -0,0 +1,54 @@ +"""CPU tests for PPO policy-loss clipping and its training-path wiring.""" + +import ast +from pathlib import Path + +import pytest +import torch + +from vime.utils.ppo_utils import compute_policy_loss + +NUM_GPUS = 0 + + +def test_compute_policy_loss_applies_dual_clip_to_negative_advantages(): + ratios = torch.tensor([2.0, 2.0, 0.5]) + ppo_kl = -ratios.log() + advantages = torch.tensor([2.0, -2.0, 2.0]) + + losses, _ = compute_policy_loss( + ppo_kl, + advantages, + eps_clip=0.2, + eps_clip_high=0.2, + eps_clip_c=1.5, + ) + + torch.testing.assert_close(losses, torch.tensor([-2.4, 3.0, -1.0])) + + +def test_policy_loss_function_forwards_eps_clip_c(): + loss_path = Path(__file__).parents[1] / "vime" / "backends" / "megatron_utils" / "loss.py" + module = ast.parse(loss_path.read_text()) + policy_loss_function = next( + node for node in module.body if isinstance(node, ast.FunctionDef) and node.name == "policy_loss_function" + ) + compute_policy_loss_calls = [ + node + for node in ast.walk(policy_loss_function) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "compute_policy_loss" + ] + + assert len(compute_policy_loss_calls) == 1 + eps_clip_c_keyword = next( + (keyword for keyword in compute_policy_loss_calls[0].keywords if keyword.arg == "eps_clip_c"), + None, + ) + assert eps_clip_c_keyword is not None + assert ast.dump(eps_clip_c_keyword.value) == ast.dump( + ast.Attribute(value=ast.Name(id="args", ctx=ast.Load()), attr="eps_clip_c", ctx=ast.Load()) + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_ppo_kl_metric.py b/tests/test_ppo_kl_metric.py new file mode 100644 index 000000000..7af2fc53d --- /dev/null +++ b/tests/test_ppo_kl_metric.py @@ -0,0 +1,68 @@ +import sys +import types +from argparse import Namespace + +import pytest +import torch + +from vime.utils.ppo_utils import compute_approx_kl + +NUM_GPUS = 0 + + +def test_ppo_estimator_does_not_corrupt_logged_kl(monkeypatch): + previous_loss = sys.modules.pop("vime.backends.megatron_utils.loss", None) + previous_cp_utils = sys.modules.pop("vime.backends.megatron_utils.cp_utils", None) + + mpu_stub = types.SimpleNamespace( + get_context_parallel_world_size=lambda: 1, + get_context_parallel_rank=lambda: 0, + is_pipeline_last_stage=lambda: True, + ) + megatron_mod = types.ModuleType("megatron") + core_mod = types.ModuleType("megatron.core") + core_mod.mpu = mpu_stub + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.core", core_mod) + + try: + from vime.backends.megatron_utils.loss import compute_advantages_and_returns + + log_probs = [torch.tensor([0.5, 0.7, 0.9])] + ref_log_probs = [torch.tensor([0.4, 0.5, 0.6])] + expected_kl = compute_approx_kl(log_probs[0], ref_log_probs[0], kl_loss_type="k1") + rollout_data = { + "log_probs": log_probs, + "ref_log_probs": ref_log_probs, + "rewards": [1.0], + "values": [torch.zeros(3)], + "response_lengths": [3], + "total_lengths": [5], + "loss_masks": [torch.ones(3)], + } + args = Namespace( + advantage_estimator="ppo", + kl_coef=0.05, + kl_loss_type="k1", + use_rollout_logprobs=False, + custom_advantage_function_path=None, + normalize_advantages=False, + use_opd=False, + gamma=1.0, + lambd=1.0, + ) + compute_advantages_and_returns(args, rollout_data) + torch.testing.assert_close(rollout_data["kl"][0], expected_kl) + finally: + if previous_loss is None: + sys.modules.pop("vime.backends.megatron_utils.loss", None) + else: + sys.modules["vime.backends.megatron_utils.loss"] = previous_loss + if previous_cp_utils is None: + sys.modules.pop("vime.backends.megatron_utils.cp_utils", None) + else: + sys.modules["vime.backends.megatron_utils.cp_utils"] = previous_cp_utils + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_ppo_logprob_entropy.py b/tests/test_ppo_logprob_entropy.py new file mode 100644 index 000000000..2299cba29 --- /dev/null +++ b/tests/test_ppo_logprob_entropy.py @@ -0,0 +1,420 @@ +"""CPU tests for fused PPO log-probability and entropy calculation.""" + +from __future__ import annotations + +import os +import socket + +import pytest +import torch + +from vime.utils.ppo_utils import calculate_log_probs_and_entropy + + +NUM_GPUS = 0 + +STRICT_ATOL = 1e-8 +STRICT_RTOL = 0.0 + + +def _free_port() -> int: + sock = socket.socket() + sock.bind(("", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def _unfused_reference_logprob_entropy( + logits: torch.Tensor, + tokens: torch.Tensor, + keep_mask: torch.Tensor | None, + *, + with_entropy: bool, + num_partitions: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Reference for the pre-fused behavior, preserving its reduction order.""" + logprob_logits = logits + if keep_mask is not None: + logprob_logits = logits.masked_fill(~keep_mask, float("-inf")) + # Match replay behavior: the sampled token must stay finite even + # when an engine-side top-p mask omitted it. + rows = torch.arange(tokens.numel(), device=logits.device) + logprob_logits[rows, tokens] = logits[rows, tokens] + + log_probs = _reference_log_probs_with_partition_order(logprob_logits, tokens, num_partitions=num_partitions) + entropy = None + if with_entropy: + entropy = _reference_entropy_with_partition_order(logits, num_partitions=num_partitions) + return log_probs, entropy + + +def _sum_in_partition_order(chunks: list[torch.Tensor]) -> torch.Tensor: + total = chunks[0] + for chunk in chunks[1:]: + total = total + chunk + return total + + +def _reference_log_probs_with_partition_order( + logits: torch.Tensor, + tokens: torch.Tensor, + *, + num_partitions: int, +) -> torch.Tensor: + rows = torch.arange(tokens.numel(), device=logits.device) + chunks = list(logits.chunk(num_partitions, dim=-1)) + vocab_per_partition = chunks[0].size(-1) + + logits_max = torch.stack([chunk.max(dim=-1, keepdim=True).values for chunk in chunks], dim=0).max(dim=0).values + normalized_chunks = [chunk - logits_max for chunk in chunks] + exp_chunks = [chunk.exp() for chunk in normalized_chunks] + sum_exp_logits = _sum_in_partition_order([chunk.sum(dim=-1, keepdim=True) for chunk in exp_chunks]) + + predicted_logits = logits.new_zeros((tokens.numel(), 1)) + for partition, normalized_chunk in enumerate(normalized_chunks): + vocab_start = partition * vocab_per_partition + local_tokens = tokens - vocab_start + on_partition = (local_tokens >= 0) & (local_tokens < vocab_per_partition) + local_tokens = local_tokens.clamp(0, vocab_per_partition - 1) + partition_predicted_logits = normalized_chunk[rows, local_tokens].unsqueeze(-1) + partition_predicted_logits = partition_predicted_logits.masked_fill(~on_partition.unsqueeze(-1), 0.0) + predicted_logits = predicted_logits + partition_predicted_logits + + return predicted_logits - sum_exp_logits.log() + + +def _reference_entropy_with_partition_order(logits: torch.Tensor, *, num_partitions: int) -> torch.Tensor: + chunks = list(logits.chunk(num_partitions, dim=-1)) + logits_max = torch.stack([chunk.max(dim=-1, keepdim=True).values for chunk in chunks], dim=0).max(dim=0).values + normalized_chunks = [chunk - logits_max for chunk in chunks] + exp_chunks = [chunk.exp() for chunk in normalized_chunks] + sum_exp_logits = _sum_in_partition_order([chunk.sum(dim=-1, keepdim=True) for chunk in exp_chunks]) + softmax_chunks = [chunk / sum_exp_logits for chunk in exp_chunks] + sum_softmax_times_logits = _sum_in_partition_order( + [(softmax * chunk).sum(dim=-1, keepdim=True) for softmax, chunk in zip(softmax_chunks, chunks, strict=True)] + ) + return (logits_max + sum_exp_logits.log() - sum_softmax_times_logits).squeeze(dim=-1) + + +def _reference_grad_with_partition_order( + logits: torch.Tensor, + tokens: torch.Tensor, + keep_mask: torch.Tensor | None, + *, + with_entropy: bool, + logprob_weights: torch.Tensor, + entropy_weights: torch.Tensor, + num_partitions: int = 1, +) -> torch.Tensor: + logprob_logits = logits + if keep_mask is not None: + logprob_logits = logits.masked_fill(~keep_mask, float("-inf")) + rows = torch.arange(tokens.numel(), device=logits.device) + logprob_logits[rows, tokens] = logits[rows, tokens] + + logprob_softmax_chunks = _reference_softmax_chunks_with_partition_order( + logprob_logits, + num_partitions=num_partitions, + ) + grad_chunks = [] + vocab_per_partition = logprob_softmax_chunks[0].size(-1) + for partition, softmax_chunk in enumerate(logprob_softmax_chunks): + vocab_start = partition * vocab_per_partition + local_tokens = tokens - vocab_start + on_partition = (local_tokens >= 0) & (local_tokens < vocab_per_partition) + local_tokens = local_tokens.clamp(0, vocab_per_partition - 1) + + grad_chunk = -softmax_chunk + rows = torch.arange(tokens.numel(), device=logits.device) + grad_2d = grad_chunk.view(-1, vocab_per_partition) + grad_2d[rows, local_tokens] += on_partition.to(dtype=grad_2d.dtype) + grad_chunk = grad_chunk * logprob_weights.reshape(-1, 1) + grad_chunks.append(grad_chunk) + + grad = torch.cat(grad_chunks, dim=-1) + + if with_entropy: + entropy_softmax_chunks = _reference_softmax_chunks_with_partition_order( + logits, + num_partitions=num_partitions, + ) + logits_chunks = list(logits.chunk(num_partitions, dim=-1)) + sum_softmax_times_logits = _sum_in_partition_order( + [ + (softmax * logits_chunk).sum(dim=-1, keepdim=True) + for softmax, logits_chunk in zip(entropy_softmax_chunks, logits_chunks, strict=True) + ] + ) + entropy_grad = torch.cat( + [ + softmax * (sum_softmax_times_logits - logits_chunk) * entropy_weights.reshape(-1, 1) + for softmax, logits_chunk in zip(entropy_softmax_chunks, logits_chunks, strict=True) + ], + dim=-1, + ) + grad = grad + entropy_grad + + return grad + + +def _reference_softmax_chunks_with_partition_order( + logits: torch.Tensor, + *, + num_partitions: int, +) -> list[torch.Tensor]: + chunks = list(logits.chunk(num_partitions, dim=-1)) + logits_max = torch.stack([chunk.max(dim=-1, keepdim=True).values for chunk in chunks], dim=0).max(dim=0).values + normalized_chunks = [chunk - logits_max for chunk in chunks] + exp_chunks = [chunk.exp() for chunk in normalized_chunks] + sum_exp_logits = _sum_in_partition_order([chunk.sum(dim=-1, keepdim=True) for chunk in exp_chunks]) + return [chunk / sum_exp_logits for chunk in exp_chunks] + + +def _single_rank_logits() -> torch.Tensor: + return torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [4.0, 1.0, 0.5, 2.0], + [-1.0, 3.0, 2.0, 0.0], + ], + dtype=torch.float32, + ) + + +def _single_rank_keep_mask() -> torch.Tensor: + return torch.tensor( + [ + [False, True, True, False], # target 3 is deliberately absent. + [False, True, False, True], # target 0 is deliberately absent. + [True, True, False, False], + ], + dtype=torch.bool, + ) + + +def _weighted_loss( + log_probs: torch.Tensor, + entropy: torch.Tensor | None, + *, + logprob_weights: torch.Tensor, + entropy_weights: torch.Tensor, +) -> torch.Tensor: + loss = (log_probs.squeeze(-1) * logprob_weights).sum() + if entropy is not None: + loss = loss + (entropy * entropy_weights).sum() + return loss + + +@pytest.mark.parametrize("chunk_size", [-1, 1, 2, 8]) +@pytest.mark.parametrize("with_mask", [False, True]) +@pytest.mark.parametrize("with_entropy", [False, True]) +def test_calculate_log_probs_and_entropy_matches_unfused_reference_single_rank( + chunk_size: int, + with_mask: bool, + with_entropy: bool, +): + logits = _single_rank_logits().requires_grad_() + tokens = torch.tensor([3, 0, 1], dtype=torch.long) + keep_mask = _single_rank_keep_mask() if with_mask else None + + log_probs, entropy = calculate_log_probs_and_entropy( + logits, + tokens, + tp_group=None, + with_entropy=with_entropy, + chunk_size=chunk_size, + log_prob_keep_mask=keep_mask, + ) + + ref_logits = logits.detach().clone().requires_grad_() + expected_log_probs, expected_entropy = _unfused_reference_logprob_entropy( + ref_logits, + tokens, + keep_mask, + with_entropy=with_entropy, + ) + + torch.testing.assert_close(log_probs, expected_log_probs, rtol=STRICT_RTOL, atol=STRICT_ATOL) + if with_entropy: + torch.testing.assert_close(entropy, expected_entropy, rtol=STRICT_RTOL, atol=STRICT_ATOL) + else: + assert entropy is None + assert expected_entropy is None + + logprob_weights = torch.tensor([0.25, -0.5, 1.5], dtype=torch.float32) + entropy_weights = torch.tensor([0.55, -0.2, 1.8], dtype=torch.float32) + loss = _weighted_loss( + log_probs, + entropy, + logprob_weights=logprob_weights, + entropy_weights=entropy_weights, + ) + loss.backward() + expected_grad = _reference_grad_with_partition_order( + ref_logits, + tokens, + keep_mask, + with_entropy=with_entropy, + logprob_weights=logprob_weights, + entropy_weights=entropy_weights, + ) + + torch.testing.assert_close(logits.grad, expected_grad, rtol=STRICT_RTOL, atol=STRICT_ATOL) + + +@pytest.mark.parametrize("with_entropy", [False, True]) +def test_calculate_log_probs_and_entropy_handles_empty_input(with_entropy: bool): + logits = torch.empty((0, 4), dtype=torch.float32, requires_grad=True) + tokens = torch.empty((0,), dtype=torch.long) + keep_mask = torch.empty((0, 4), dtype=torch.bool) + + log_probs, entropy = calculate_log_probs_and_entropy( + logits, + tokens, + tp_group=None, + with_entropy=with_entropy, + chunk_size=2, + log_prob_keep_mask=keep_mask, + ) + + assert log_probs.shape == (0,) + if with_entropy: + assert entropy is not None + assert entropy.shape == (0,) + else: + assert entropy is None + + +def _distributed_full_logits() -> torch.Tensor: + return torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 0.5, -1.0], + [4.0, 1.0, 0.5, 2.0, 3.0, 0.0], + [-1.0, 3.0, 2.0, 0.0, 1.0, 5.0], + [0.2, -0.4, 1.7, -2.0, 3.3, 0.0], + ], + dtype=torch.float32, + ) + + +def _distributed_keep_mask() -> torch.Tensor: + return torch.tensor( + [ + [False, True, False, True, False, False], # target 5 is absent. + [False, True, False, False, True, False], # target 0 is absent. + [True, False, True, False, False, True], # target 3 is absent. + [False, True, False, True, False, False], # target 2 is absent. + ], + dtype=torch.bool, + ) + + +def _distributed_vocab_worker( + rank: int, + world_size: int, + with_mask: bool, + with_entropy: bool, + chunk_size: int, + master_port: int, +) -> None: + import torch.distributed as dist + + torch.set_num_threads(1) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + try: + full_logits = _distributed_full_logits() + tokens = torch.tensor([5, 0, 3, 2], dtype=torch.long) + full_keep_mask = _distributed_keep_mask() if with_mask else None + + vocab_per_rank = full_logits.size(-1) // world_size + vocab_start = rank * vocab_per_rank + vocab_end = vocab_start + vocab_per_rank + local_logits = full_logits[:, vocab_start:vocab_end].detach().clone().requires_grad_() + local_keep_mask = None + if full_keep_mask is not None: + local_keep_mask = full_keep_mask[:, vocab_start:vocab_end] + + log_probs, entropy = calculate_log_probs_and_entropy( + local_logits, + tokens, + tp_group=None, + with_entropy=with_entropy, + chunk_size=chunk_size, + log_prob_keep_mask=local_keep_mask, + ) + + ref_logits = full_logits.detach().clone().requires_grad_() + expected_log_probs, expected_entropy = _unfused_reference_logprob_entropy( + ref_logits, + tokens, + full_keep_mask, + with_entropy=with_entropy, + num_partitions=world_size, + ) + + torch.testing.assert_close(log_probs, expected_log_probs, rtol=STRICT_RTOL, atol=STRICT_ATOL) + if with_entropy: + torch.testing.assert_close(entropy, expected_entropy, rtol=STRICT_RTOL, atol=STRICT_ATOL) + else: + assert entropy is None + assert expected_entropy is None + + logprob_weights = torch.tensor([0.25, -0.5, 1.5, -0.75], dtype=torch.float32) + entropy_weights = torch.tensor([0.55, -0.2, 1.8, 0.4], dtype=torch.float32) + loss = _weighted_loss( + log_probs, + entropy, + logprob_weights=logprob_weights, + entropy_weights=entropy_weights, + ) + loss.backward() + expected_grad = _reference_grad_with_partition_order( + ref_logits, + tokens, + full_keep_mask, + with_entropy=with_entropy, + logprob_weights=logprob_weights, + entropy_weights=entropy_weights, + num_partitions=world_size, + ) + + torch.testing.assert_close( + local_logits.grad, + expected_grad[:, vocab_start:vocab_end], + rtol=STRICT_RTOL, + atol=STRICT_ATOL, + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "with_mask,with_entropy,chunk_size", + [ + pytest.param(False, True, -1, id="unmasked_entropy_no_chunks"), + pytest.param(True, True, 2, id="masked_entropy_chunks"), + pytest.param(True, False, -1, id="masked_logprob_only_no_chunks"), + pytest.param(False, False, 2, id="unmasked_logprob_only_chunks"), + ], +) +def test_calculate_log_probs_and_entropy_matches_unfused_reference_vocab_parallel( + with_mask: bool, + with_entropy: bool, + chunk_size: int, +): + import torch.multiprocessing as mp + + world_size = 2 + mp.spawn( + _distributed_vocab_worker, + args=(world_size, with_mask, with_entropy, chunk_size, _free_port()), + nprocs=world_size, + join=True, + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_ppo_logprob_entropy_gpu.py b/tests/test_ppo_logprob_entropy_gpu.py new file mode 100644 index 000000000..95422d067 --- /dev/null +++ b/tests/test_ppo_logprob_entropy_gpu.py @@ -0,0 +1,355 @@ +"""CUDA parity test for fused PPO log-probability and entropy calculation.""" + +from __future__ import annotations + +import os +import socket + +import pytest +import torch + +from vime.utils.ppo_utils import calculate_log_probs_and_entropy + + +NUM_GPUS = 2 + +# Megatron's JIT fused CE can differ from the same Python-level expression by +# one fp32 ulp in the unmasked path. +FORWARD_ATOL = 1e-7 +FORWARD_RTOL = 0.0 +# Entropy values are O(1) in this parity fixture; allow a small difference from +# the memory-saving CUDA reduction without relaxing log-prob parity. +ENTROPY_FORWARD_ATOL = 1e-4 +BACKWARD_ATOL = 1e-8 +BACKWARD_RTOL = 0.0 +# Entropy backward uses a separate memory-saving CUDA reduction. +ENTROPY_BACKWARD_ATOL = 1e-6 + +PARITY_SCENARIOS = [ + (-1, False, False, False), + (-1, False, True, False), + (-1, False, True, True), + (-1, True, False, False), + (-1, True, True, False), + (-1, True, True, True), + (2, False, False, False), + (2, False, True, False), + (2, False, True, True), + (2, True, False, False), + (2, True, True, False), + (2, True, True, True), +] + + +def _free_port() -> int: + sock = socket.socket() + sock.bind(("", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def _full_logits() -> torch.Tensor: + return torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 0.5, -1.0], + [4.0, 1.0, 0.5, 2.0, 3.0, 0.0], + [-1.0, 3.0, 2.0, 0.0, 1.0, 5.0], + [0.2, -0.4, 1.7, -2.0, 3.3, 0.0], + ], + dtype=torch.float32, + ) + + +def _keep_mask() -> torch.Tensor: + return torch.tensor( + [ + [False, True, False, True, False, False], # target 5 is absent. + [False, True, False, False, True, False], # target 0 is absent. + [True, False, True, False, False, True], # target 3 is absent. + [False, True, False, True, False, False], # target 2 is absent. + ], + dtype=torch.bool, + ) + + +def _weighted_loss( + log_probs: torch.Tensor, + entropy: torch.Tensor | None, + *, + logprob_weights: torch.Tensor, + entropy_weights: torch.Tensor | None, +) -> torch.Tensor: + loss = (log_probs.squeeze(-1) * logprob_weights).sum() + if entropy is not None and entropy_weights is not None: + loss = loss + (entropy * entropy_weights).sum() + return loss + + +def _legacy_compute_log_probs( + logits: torch.Tensor, + tokens: torch.Tensor, + process_group, + keep_mask: torch.Tensor | None = None, +) -> torch.Tensor: + from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy + + if keep_mask is not None: + keep_mask = keep_mask.clone() + vocab_local = keep_mask.size(-1) + vocab_start = process_group.rank() * vocab_local + local_tokens = tokens - vocab_start + on_shard = (local_tokens >= 0) & (local_tokens < vocab_local) + rows = torch.nonzero(on_shard, as_tuple=False).squeeze(-1) + if rows.numel() > 0: + keep_mask[rows, local_tokens[rows]] = True + logits = logits.masked_fill(~keep_mask, float("-inf")) + + return -fused_vocab_parallel_cross_entropy(logits.unsqueeze(1), tokens.unsqueeze(1), process_group) + + +class _LegacyVocabParallelEntropy(torch.autograd.Function): + @staticmethod + def forward(ctx, vocab_parallel_logits: torch.Tensor, process_group) -> torch.Tensor: + logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values + torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=process_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) + torch.distributed.all_reduce(normalized_sum_exp_logits, group=process_group) + softmax_logits = normalized_exp_logits.div_(normalized_sum_exp_logits) + sum_softmax_times_logits = (softmax_logits * vocab_parallel_logits).sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(sum_softmax_times_logits, group=process_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) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors + grad_input = softmax_logits * (sum_softmax_times_logits - vocab_parallel_logits) + grad_input = grad_input * grad_output.unsqueeze(dim=-1) + return grad_input, None + + +def _legacy_compute_entropy_from_logits(logits: torch.Tensor, process_group) -> torch.Tensor: + return _LegacyVocabParallelEntropy.apply(logits, process_group) + + +def _assert_logprob_backward_close(actual_grad: torch.Tensor, legacy_grad: torch.Tensor) -> None: + # Megatron's fused vocab-parallel CE backward quantizes the log-prob + # gradient to bfloat16 on CUDA. The new implementation keeps fp32 grads, so + # compare this branch at the legacy kernel's effective precision. + if actual_grad.is_cuda: + actual_grad = actual_grad.to(torch.bfloat16) + legacy_grad = legacy_grad.to(torch.bfloat16) + torch.testing.assert_close(actual_grad, legacy_grad, rtol=BACKWARD_RTOL, atol=BACKWARD_ATOL) + + +def _assert_legacy_parity( + *, + process_group, + device: torch.device, + logits: torch.Tensor, + tokens: torch.Tensor, + keep_mask: torch.Tensor | None, + chunk_size: int, + with_entropy: bool, + entropy_has_grad: bool, +) -> None: + log_probs, entropy = calculate_log_probs_and_entropy( + logits, + tokens, + tp_group=process_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + log_prob_keep_mask=keep_mask, + with_entropy_grad=entropy_has_grad, + ) + + legacy_logits = logits.detach().clone().requires_grad_() + legacy_log_probs = _legacy_compute_log_probs(legacy_logits.clone(), tokens, process_group, keep_mask=keep_mask) + + torch.testing.assert_close(log_probs, legacy_log_probs, rtol=FORWARD_RTOL, atol=FORWARD_ATOL) + if with_entropy: + legacy_entropy = _legacy_compute_entropy_from_logits(legacy_logits.clone(), process_group) + torch.testing.assert_close(entropy, legacy_entropy, rtol=FORWARD_RTOL, atol=ENTROPY_FORWARD_ATOL) + assert entropy.requires_grad == entropy_has_grad + else: + legacy_entropy = None + assert entropy is None + + logprob_weights = torch.tensor([0.25, -0.5, 1.5, -0.75], dtype=torch.float32, device=device) + logprob_logits = logits.detach().clone().requires_grad_() + logprob_values, _ = calculate_log_probs_and_entropy( + logprob_logits, + tokens, + tp_group=process_group, + with_entropy=with_entropy, + chunk_size=chunk_size, + log_prob_keep_mask=keep_mask, + with_entropy_grad=entropy_has_grad, + ) + legacy_logprob_logits = logits.detach().clone().requires_grad_() + legacy_logprob_values = _legacy_compute_log_probs( + legacy_logprob_logits.clone(), tokens, process_group, keep_mask=keep_mask + ) + _weighted_loss( + logprob_values, + None, + logprob_weights=logprob_weights, + entropy_weights=None, + ).backward() + _weighted_loss( + legacy_logprob_values, + None, + logprob_weights=logprob_weights, + entropy_weights=None, + ).backward() + _assert_logprob_backward_close(logprob_logits.grad, legacy_logprob_logits.grad) + + if with_entropy and entropy_has_grad: + entropy_weights = torch.tensor([0.55, -0.2, 1.8, 0.4], dtype=torch.float32, device=device) + entropy_logits = logits.detach().clone().requires_grad_() + _, entropy_values = calculate_log_probs_and_entropy( + entropy_logits, + tokens, + tp_group=process_group, + with_entropy=True, + chunk_size=chunk_size, + log_prob_keep_mask=keep_mask, + with_entropy_grad=True, + ) + legacy_entropy_logits = logits.detach().clone().requires_grad_() + legacy_entropy_values = _legacy_compute_entropy_from_logits(legacy_entropy_logits.clone(), process_group) + (entropy_values * entropy_weights).sum().backward() + (legacy_entropy_values * entropy_weights).sum().backward() + torch.testing.assert_close( + entropy_logits.grad, + legacy_entropy_logits.grad, + rtol=BACKWARD_RTOL, + atol=ENTROPY_BACKWARD_ATOL, + ) + + +@pytest.fixture(scope="module") +def nccl_process_group(): + import torch.distributed as dist + + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + pytest.importorskip("megatron.core.fusions.fused_cross_entropy") + if not dist.is_nccl_available(): + pytest.skip("NCCL is required") + + created_process_group = False + if dist.is_initialized(): + process_group = dist.group.WORLD + if dist.get_backend(process_group) != "nccl": + pytest.skip("legacy Megatron CUDA parity needs an NCCL process group") + else: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(_free_port()) + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created_process_group = True + process_group = dist.group.WORLD + + yield process_group + + if created_process_group: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "with_entropy,entropy_has_grad", + [ + pytest.param(False, False, id="without_entropy"), + pytest.param(True, False, id="entropy_forward_only"), + pytest.param(True, True, id="entropy_backward"), + ], +) +@pytest.mark.parametrize("with_mask", [False, True], ids=["unmasked", "masked"]) +@pytest.mark.parametrize("chunk_size", [-1, 2], ids=["no_chunks", "chunks"]) +def test_calculate_log_probs_and_entropy_matches_legacy_megatron_cuda( + nccl_process_group, + chunk_size: int, + with_mask: bool, + with_entropy: bool, + entropy_has_grad: bool, +): + process_group = nccl_process_group + torch.cuda.set_device(0) + device = torch.device("cuda", torch.cuda.current_device()) + logits = _full_logits().to(device=device).requires_grad_() + tokens = torch.tensor([5, 0, 3, 2], dtype=torch.long, device=device) + keep_mask = _keep_mask().to(device=device) if with_mask else None + + _assert_legacy_parity( + process_group=process_group, + device=device, + logits=logits, + tokens=tokens, + keep_mask=keep_mask, + chunk_size=chunk_size, + with_entropy=with_entropy, + entropy_has_grad=entropy_has_grad, + ) + + +def _tp2_worker(rank: int, world_size: int, master_port: int) -> None: + import torch.distributed as dist + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + torch.cuda.set_device(rank) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + try: + process_group = dist.group.WORLD + device = torch.device("cuda", rank) + full_logits = _full_logits().to(device=device) + full_keep_mask = _keep_mask().to(device=device) + tokens = torch.tensor([5, 0, 3, 2], dtype=torch.long, device=device) + + vocab_per_rank = full_logits.size(-1) // world_size + vocab_start = rank * vocab_per_rank + vocab_end = vocab_start + vocab_per_rank + for chunk_size, with_mask, with_entropy, entropy_has_grad in PARITY_SCENARIOS: + logits = full_logits[:, vocab_start:vocab_end].detach().clone().requires_grad_() + keep_mask = full_keep_mask[:, vocab_start:vocab_end] if with_mask else None + _assert_legacy_parity( + process_group=process_group, + device=device, + logits=logits, + tokens=tokens, + keep_mask=keep_mask, + chunk_size=chunk_size, + with_entropy=with_entropy, + entropy_has_grad=entropy_has_grad, + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_calculate_log_probs_and_entropy_matches_legacy_megatron_cuda_tp2(): + pytest.importorskip("megatron.core.fusions.fused_cross_entropy") + if torch.cuda.device_count() < 2: + pytest.skip("TP=2 parity requires two CUDA devices") + + import torch.distributed as dist + import torch.multiprocessing as mp + + if not dist.is_nccl_available(): + pytest.skip("NCCL is required") + + world_size = 2 + mp.spawn( + _tp2_worker, + args=(world_size, _free_port()), + nprocs=world_size, + join=True, + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_process_rollout_data.py b/tests/test_process_rollout_data.py new file mode 100644 index 000000000..06f664272 --- /dev/null +++ b/tests/test_process_rollout_data.py @@ -0,0 +1,166 @@ +"""CPU unit tests for ``vime.utils.data.process_rollout_data``. + +``RolloutManager._split_train_data_by_dp`` ships two kinds of fields to the +trainer: + + * per-sample fields (``response_lengths``, ``loss_masks``, ...) already + sliced down to the samples this DP rank owns, and + * ``raw_reward`` / ``total_lengths``, sent whole because something on the + training side still needs the full rollout batch. + +``process_rollout_data`` is where the second group is reconciled with the +first. These tests pin that contract: + + 1. ``total_lengths`` comes out DP-local (already the case). + 2. ``raw_reward`` stays global — ``log_passrate`` reshapes it into + ``[rollout_batch_size, n_samples_per_prompt]`` groups, which only works + on the full batch. + 3. ``local_raw_reward`` is the DP-local view, positionally aligned with the + per-sample fields. + +(3) is the regression guard for the ``--log-correct-samples`` crash: that +block zips rewards against ``response_lengths`` / ``total_lengths`` / +``loss_masks`` / ``log_probs`` by position, so feeding it the global +``raw_reward`` walked off the end of this rank's lists with an +``IndexError`` (and, before running off the end, silently attributed one +sample's reward to a different sample). +""" + +from __future__ import annotations + +import pytest +import ray + +from vime.utils.data import process_rollout_data + +NUM_GPUS = 0 + + +class _FakeBox: + """Stand-in for ``vime.ray.utils.Box``: payload lives behind ``.inner``.""" + + def __init__(self, inner): + self.inner = inner + + +@pytest.fixture +def unwrap_ray_get(monkeypatch): + """``process_rollout_data`` uses Ray only to deref the per-rank Box. + + Patching ``ray.get`` to the identity keeps these tests single-process + (no cluster start-up) while still exercising the real function. + """ + monkeypatch.setattr(ray, "get", lambda ref: ref) + + +def _split_train_data_by_dp(partitions, raw_reward, response_lengths, total_lengths): + """Mirror what ``RolloutManager._split_train_data_by_dp`` packages per rank.""" + return [ + _FakeBox( + { + "partition": partition, + "response_lengths": [response_lengths[j] for j in partition], + "raw_reward": list(raw_reward), + "total_lengths": list(total_lengths), + } + ) + for partition in partitions + ] + + +# 8 samples; only the odd-indexed ones are correct. Lengths encode their own +# global index so a mis-pairing is visible in the assertion message. +RAW_REWARD = [0, 1, 0, 1, 0, 1, 0, 1] +RESPONSE_LENGTHS = [100, 101, 102, 103, 104, 105, 106, 107] +TOTAL_LENGTHS = [200, 201, 202, 203, 204, 205, 206, 207] + + +@pytest.mark.parametrize( + "partitions", + [ + pytest.param([[0, 2, 4, 6], [1, 3, 5, 7]], id="dp2-interleaved"), + pytest.param([[0, 1, 2, 3], [4, 5, 6, 7]], id="dp2-contiguous"), + pytest.param([[0, 3], [1, 6], [2, 5], [4, 7]], id="dp4-balanced"), + # Even at dp_size=1 the partition is a permutation: first-fit packing + # reorders samples by length. + pytest.param([[3, 0, 7, 1, 5, 2, 6, 4]], id="dp1-permuted"), + ], +) +def test_local_raw_reward_is_dp_local_and_aligned(unwrap_ray_get, partitions): + dp_size = len(partitions) + refs = _split_train_data_by_dp(partitions, RAW_REWARD, RESPONSE_LENGTHS, TOTAL_LENGTHS) + + for dp_rank, partition in enumerate(partitions): + rollout_data = process_rollout_data(rollout_data_ref=refs, dp_rank=dp_rank, dp_size=dp_size) + + local_raw_reward = rollout_data["local_raw_reward"] + assert local_raw_reward == [RAW_REWARD[j] for j in partition] + # Positional alignment with the per-sample fields is the whole point. + assert len(local_raw_reward) == len(rollout_data["response_lengths"]) + assert len(local_raw_reward) == len(rollout_data["total_lengths"]) + + +@pytest.mark.parametrize( + "partitions", + [ + pytest.param([[0, 2, 4, 6], [1, 3, 5, 7]], id="dp2-interleaved"), + pytest.param([[3, 0, 7, 1, 5, 2, 6, 4]], id="dp1-permuted"), + ], +) +def test_correct_sample_selection_matches_owned_samples(unwrap_ray_get, partitions): + """Replay the ``--log-correct-samples`` selection loop. + + Regression for the ``IndexError`` this used to raise on DP > 1. + """ + dp_size = len(partitions) + refs = _split_train_data_by_dp(partitions, RAW_REWARD, RESPONSE_LENGTHS, TOTAL_LENGTHS) + + for dp_rank, partition in enumerate(partitions): + rollout_data = process_rollout_data(rollout_data_ref=refs, dp_rank=dp_rank, dp_size=dp_size) + + response_lengths = rollout_data["response_lengths"] + total_lengths = rollout_data["total_lengths"] + + correct_response_lengths = [] + correct_total_lengths = [] + for i, raw_reward in enumerate(rollout_data["local_raw_reward"]): + if raw_reward == 1: + correct_response_lengths.append(response_lengths[i]) + correct_total_lengths.append(total_lengths[i]) + + expected = [j for j in partition if RAW_REWARD[j] == 1] + assert correct_response_lengths == [RESPONSE_LENGTHS[j] for j in expected] + assert correct_total_lengths == [TOTAL_LENGTHS[j] for j in expected] + + +def test_raw_reward_stays_global(unwrap_ray_get): + """``log_passrate`` needs the whole batch, so the global copy must survive.""" + partitions = [[0, 2, 4, 6], [1, 3, 5, 7]] + refs = _split_train_data_by_dp(partitions, RAW_REWARD, RESPONSE_LENGTHS, TOTAL_LENGTHS) + + for dp_rank in range(len(partitions)): + rollout_data = process_rollout_data(rollout_data_ref=refs, dp_rank=dp_rank, dp_size=len(partitions)) + assert rollout_data["raw_reward"] == RAW_REWARD + + +def test_missing_raw_reward_is_tolerated(unwrap_ray_get): + """Forward-only passes ship no ``raw_reward``; don't invent one.""" + partition = [1, 0] + refs = [ + _FakeBox( + { + "partition": partition, + "response_lengths": [101, 100], + "total_lengths": [200, 201], + } + ) + ] + + rollout_data = process_rollout_data(rollout_data_ref=refs, dp_rank=0, dp_size=1) + + assert "local_raw_reward" not in rollout_data + assert rollout_data["total_lengths"] == [201, 200] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py b/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py index c660f8f34..53df35172 100644 --- a/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py +++ b/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py @@ -77,7 +77,6 @@ def _common_args(debug_data_dir: str): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 8 " "--colocate " - "--megatron-to-hf-mode bridge " ) return f"{ckpt_args} " f"{rollout_args} " f"{optimizer_args} " f"{grpo_args} " f"{perf_args} " f"{misc_args} " diff --git a/tests/test_qwen2.5_0.5B_debug_train_dump_e2e.py b/tests/test_qwen2.5_0.5B_debug_train_dump_e2e.py new file mode 100644 index 000000000..1e0f83d04 --- /dev/null +++ b/tests/test_qwen2.5_0.5B_debug_train_dump_e2e.py @@ -0,0 +1,169 @@ +"""End-to-end debug-dump alignment test with TP/PP/CP all enabled. + +Runs a single rollout+train pass on Qwen2.5-0.5B with tensor / pipeline / +context parallel all > 1 and ``--dump-details``, which writes both the rollout +debug dump (``rollout_data/{id}.pt``) and the train debug dump +(``train_data/{id}.pt``). It then joins the two dumps by ``rollout_position`` +and asserts that the train dump's CP-reassembled ``rollout_log_probs`` match the +per-sample ``rollout_log_probs`` stored on the rollout side. + +This is the strongest correctness check for the train dump: it exercises the +writer selection (only the last PP stage + TP rank 0 write), the cross-CP +reassembly of response-token fields, and the sample-index/position ordering all +at once — a mismatch in any of them makes the compared log-probs diverge. + +Uses Qwen2.5-0.5B-Instruct with 8 GPUs: TP=2, PP=2, CP=2 (DP=1). +""" + +import os +import tempfile + +import torch + +import vime.utils.external_utils.command_utils as U + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 8 +NUM_ROLLOUT = 1 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + + +def _train_args(dump_dir: str) -> str: + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ --ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + f"--num-rollout {NUM_ROLLOUT} " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 256 " + "--rollout-temperature 0.8 " + "--global-batch-size 16 " + ) + + # TP=2, PP=2, CP=2 -> 8 GPUs, DP=1. Exercises the dump's writer selection + # (last PP stage + TP0 + CP0) and the cross-CP response-field reassembly. + parallel_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 2 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + grpo_args = "--advantage-estimator grpo --eps-clip 0.2 " + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--ci-test " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 2 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-cudagraph-capture-size 16 " + ) + + return ( + f"{ckpt_args} {rollout_args} {optimizer_args} {grpo_args} " + f"{parallel_args} {misc_args} {vllm_args} " + f"--dump-details {dump_dir} " + ) + + +def _verify(dump_dir: str): + """Join the rollout and train dumps by position and compare rollout_log_probs.""" + for rollout_id in range(NUM_ROLLOUT): + rollout_path = f"{dump_dir}/rollout_data/{rollout_id}.pt" + train_path = f"{dump_dir}/train_data/{rollout_id}.pt" + + rollout = torch.load(rollout_path, weights_only=False) + train = torch.load(train_path, weights_only=True) + assert train["format_version"] == 2, train["format_version"] + + rollout_samples = rollout["samples"] + train_samples = train["samples"] + assert train_samples, f"empty train dump for rollout {rollout_id}" + + compared = 0 + for sample in train_samples: + pos = sample["rollout_position"] + assert pos is not None, "rollout_position missing; cannot align train dump to rollout dump" + train_lp = sample.get("rollout_log_probs") + if train_lp is None: + continue # rollout log-probs not carried into training; nothing to compare + ref_lp = rollout_samples[pos]["rollout_log_probs"] + assert ref_lp is not None, f"rollout dump sample {pos} has no rollout_log_probs" + ref_lp = torch.as_tensor(ref_lp, dtype=train_lp.dtype) + assert train_lp.shape == ref_lp.shape, ( + f"rollout {rollout_id} pos {pos}: reassembled shape {tuple(train_lp.shape)} " + f"!= rollout dump {tuple(ref_lp.shape)}" + ) + torch.testing.assert_close( + train_lp, + ref_lp, + rtol=1e-3, + atol=1e-3, + msg=lambda m, pos=pos, rid=rollout_id: f"rollout {rid} pos {pos}: rollout_log_probs mismatch\n{m}", + ) + compared += 1 + + assert compared > 0, ( + f"rollout {rollout_id}: no rollout_log_probs were compared; the train dump did not carry " + "rollout_log_probs, so this test would silently pass without checking anything." + ) + print( + f"rollout {rollout_id}: verified {compared} samples' reassembled rollout_log_probs match the rollout dump" + ) + + +def execute(): + dump_dir = tempfile.mkdtemp(prefix="vime_dump_details_") + print(f"Using dump-details dir: {dump_dir}") + + U.execute_train( + train_args=_train_args(dump_dir), + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + print("=" * 60) + print("Verifying train dump aligns with rollout dump (join by rollout_position)") + print("=" * 60) + _verify(dump_dir) + print("Train/rollout debug-dump alignment verified.") + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py index 915379ff5..2609d0c2b 100644 --- a/tests/test_qwen2.5_0.5B_fanout_short.py +++ b/tests/test_qwen2.5_0.5B_fanout_short.py @@ -10,7 +10,7 @@ this test, **no e2e training run had ever exercised the full chain**: custom_generate returns list[Sample] sharing rollout_id - → _validate_rollout_id_annotated at depth ≥ 2 passes + → validate_rollout_id_annotated at depth ≥ 2 passes → _split_train_data_by_dp groups by rollout_id and trims to N steps using ``rollout_batch_size * n_samples_per_prompt / global_batch_size`` (NOT total sample count, which would inflate steps once N>1) @@ -20,10 +20,9 @@ num_rollouts (not num_samples), keeping grad magnitude stable independent of fan-out -The fan-out function itself lives in -``vime/rollout/_fanout_test_helpers.py`` — it has to be at a dot-free -module path so ``importlib.import_module`` can resolve the string -``--custom-generate-function-path`` flag (this filename has dots). +The fan-out functions live in ``tests/fanout_test_helpers.py``. The +dedicated helper module gives ``importlib.import_module`` a valid module +path even though this E2E test's filename itself contains dots. Test choices ------------ @@ -48,6 +47,7 @@ MODEL_NAME = "Qwen2.5-0.5B-Instruct" MODEL_TYPE = "qwen2.5-0.5B" NUM_GPUS = 4 +TESTS_DIR = os.path.dirname(os.path.abspath(__file__)) # Counter file used by the compact_generate helper. We pass its path # through to the Ray-submitted job via an env var so all worker @@ -100,7 +100,7 @@ def execute(): "--rollout-temperature 0.8 " "--global-batch-size 4 " "--balance-data " - "--custom-generate-function-path vime.rollout._fanout_test_helpers.compact_generate " + "--custom-generate-function-path fanout_test_helpers.compact_generate " # GRPO normalization needs per-prompt grouping. The default # ``_post_process_rewards`` (vime/ray/rollout.py) reshapes # by ``n_samples_per_prompt`` and falls back to "one big group" @@ -110,7 +110,7 @@ def execute(): # compact_generate preserves it across siblings) so each prompt's # siblings normalize against each other, matching the GRPO # semantics the default targets in the uniform case. - "--custom-reward-post-process-path vime.rollout._fanout_test_helpers.grpo_normalize_by_group_index " + "--custom-reward-post-process-path fanout_test_helpers.grpo_normalize_by_group_index " ) perf_args = ( @@ -165,7 +165,6 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 4 " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( @@ -185,9 +184,13 @@ def execute(): train_args=train_args, num_gpus_per_node=NUM_GPUS, megatron_model_type=MODEL_TYPE, - # Make the counter path visible inside the Ray-submitted job - # (helper picks it up via os.environ). - extra_env_vars={"VIME_FANOUT_TEST_COUNTER_FILE": FANOUT_COUNTER_FILE}, + extra_env_vars={ + # Make the helper importable by both the Ray driver and workers + # without installing test modules as part of the vime package. + "PYTHONPATH": f"{TESTS_DIR}:{U.repo_base_dir}:/root/Megatron-LM/", + # The helper picks up the shared counter path via os.environ. + "VIME_FANOUT_TEST_COUNTER_FILE": FANOUT_COUNTER_FILE, + }, ) # Post-train assertion: compact_generate must have been called exactly @@ -214,8 +217,6 @@ def execute(): if __name__ == "__main__": prepare() - os.environ.pop("http_proxy") - os.environ.pop("https_proxy") - os.environ.pop("HTTP_PROXY") - os.environ.pop("HTTPS_PROXY") + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) execute() diff --git a/tests/test_qwen2.5_0.5B_fully_async_short.py b/tests/test_qwen2.5_0.5B_fully_async_short.py index 776cd29e4..92b161c5a 100644 --- a/tests/test_qwen2.5_0.5B_fully_async_short.py +++ b/tests/test_qwen2.5_0.5B_fully_async_short.py @@ -1,7 +1,6 @@ """CI smoke test for the fully-async rollout path. -Mirrors ``test_qwen2.5_0.5B_async_short`` (Qwen2.5-0.5B + dapo-math-17k + -3 rollouts of GRPO) but flips the rollout function over to +Uses Qwen2.5-0.5B with dapo-math-17k and GRPO, selecting ``vime.rollout.fully_async_rollout.generate_rollout_fully_async`` so the fully-async worker path gets exercised end-to-end. @@ -10,6 +9,9 @@ """ import os + +import torch + import vime.utils.external_utils.command_utils as U @@ -22,14 +24,25 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/dapo-math-17k") + if torch.version.hip is not None: + # ROCm image has no modelopt bridge: convert HF->Megatron into a container-local dir. + U.convert_checkpoint( + MODEL_NAME, + MODEL_TYPE, + num_gpus_per_node=1, + extra_args="--no-gradient-accumulation-fusion --attention-backend flash", + dir_dst="/tmp", + ) def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + if torch.version.hip is not None: + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ --ref-load /tmp/{MODEL_NAME}_torch_dist/ " + else: + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " rollout_args = ( - # The only line that differs from test_qwen2.5_0.5B_async_short.py: - # use the public fully-async rollout function. + # Select the public fully-async rollout function. "--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async " "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " "--input-key prompt " @@ -100,7 +113,7 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 1 " "--rollout-num-gpus 3 " - "--megatron-to-hf-mode bridge " + f'{"--no-gradient-accumulation-fusion --no-offload-train " if torch.version.hip is not None else ""}' ) train_args = ( diff --git a/tests/test_qwen2.5_0.5B_opd_vllm.py b/tests/test_qwen2.5_0.5B_opd_vllm.py index e76c0a45d..d08feafc8 100644 --- a/tests/test_qwen2.5_0.5B_opd_vllm.py +++ b/tests/test_qwen2.5_0.5B_opd_vllm.py @@ -171,7 +171,6 @@ def launch_teacher(): "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_TRAIN_GPUS} " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( diff --git a/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py b/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py deleted file mode 100644 index 140c15f8d..000000000 --- a/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py +++ /dev/null @@ -1,128 +0,0 @@ -import os -import tempfile - -import vime.utils.external_utils.command_utils as U - -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" -NUM_GPUS = 4 - - -def prepare(): - U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - U.hf_download_dataset("zhuzilin/dapo-math-17k") - - -def execute(): - megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) - megatron_config.write( - """ -megatron: - - name: default - role: critic - overrides: - lr: 1e-5 - - name: default - role: actor - overrides: - lr: 1e-6 -""" - ) - megatron_config.close() - - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " - - rollout_args = ( - "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " - "--input-key prompt " - "--label-key label " - "--apply-chat-template " - "--rollout-shuffle " - "--rm-type deepscaler " - "--num-rollout 2 " - "--rollout-batch-size 4 " - "--n-samples-per-prompt 4 " - "--rollout-max-response-len 1024 " - "--rollout-temperature 0.8 " - "--global-batch-size 16 " - "--balance-data " - ) - - perf_args = ( - "--tensor-model-parallel-size 1 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 1 " - "--expert-model-parallel-size 1 " - "--expert-tensor-parallel-size 1 " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 9216 " - ) - - ppo_args = ( - "--advantage-estimator ppo " - "--kl-loss-coef 0.00 " - "--kl-loss-type k1 " - "--kl-coef 0.00 " - "--entropy-coef 0.00 " - "--eps-clip 4e-4 " - "--num-critic-only-steps 2 " - "--normalize-advantages " - ) - - optimizer_args = ( - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - ) - - vllm_args = ( - "--rollout-num-gpus-per-engine 1 " - "--rollout-num-gpus 2 " - "--vllm-gpu-memory-utilization 0.65 " - "--vllm-max-cudagraph-capture-size 64" - ) - - ci_args = "--ci-test " - - misc_args = ( - "--attention-dropout 0.0 " - "--hidden-dropout 0.0 " - "--accumulate-allreduce-grads-in-fp32 " - "--attention-softmax-in-fp32 " - "--attention-backend flash " - "--actor-num-nodes 1 " - "--actor-num-gpus-per-node 4 " - "--megatron-to-hf-mode bridge " - "--colocate " - ) - - train_args = ( - f"--megatron-config-path {megatron_config.name} " - f"{ckpt_args} " - f"{rollout_args} " - f"{optimizer_args} " - f"{ppo_args} " - f"{U.get_default_wandb_args(__file__)} " - f"{perf_args} " - f"{vllm_args} " - f"{ci_args} " - f"{misc_args} " - ) - - U.execute_train( - train_args=train_args, - num_gpus_per_node=NUM_GPUS, - megatron_model_type=MODEL_TYPE, - ) - - -if __name__ == "__main__": - prepare() - for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): - os.environ.pop(proxy_var, None) - execute() diff --git a/tests/test_qwen2.5_0.5B_vllm_config.py b/tests/test_qwen2.5_0.5B_vllm_config.py index 878e185d8..b49fb32b0 100644 --- a/tests/test_qwen2.5_0.5B_vllm_config.py +++ b/tests/test_qwen2.5_0.5B_vllm_config.py @@ -100,7 +100,7 @@ def execute(): "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-num-seqs 32 " - "--vllm-max-cudagraph-capture-size 16 " + "--vllm-max-cudagraph-capture-size 32 " f"--vllm-config {config_path} " ) @@ -115,7 +115,6 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 8 " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( diff --git a/tests/test_qwen2.5_0.5B_vllm_config_distributed.py b/tests/test_qwen2.5_0.5B_vllm_config_distributed.py index cd63c0eed..45fe6139c 100644 --- a/tests/test_qwen2.5_0.5B_vllm_config_distributed.py +++ b/tests/test_qwen2.5_0.5B_vllm_config_distributed.py @@ -101,7 +101,7 @@ def execute(): "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-num-seqs 32 " - "--vllm-max-cudagraph-capture-size 16 " + "--vllm-max-cudagraph-capture-size 32 " f"--vllm-config {config_path} " ) @@ -117,7 +117,6 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 4 " "--rollout-num-gpus 4 " - "--megatron-to-hf-mode bridge " ) train_args = ( diff --git a/tests/test_qwen2.5_vl_3B_ep_disaggregation.py b/tests/test_qwen2.5_vl_3B_ep_disaggregation.py new file mode 100644 index 000000000..8472fb0b2 --- /dev/null +++ b/tests/test_qwen2.5_vl_3B_ep_disaggregation.py @@ -0,0 +1,292 @@ +"""Three-GPU vLLM-only acceptance test for encoder-prefill disaggregation.""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import os +import shutil +import sys +import tempfile +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import ray +import requests +import yaml +from PIL import Image + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import vime.utils.external_utils.command_utils as U +from vime.backends.vllm_utils.arguments import vllm_parse_args +from vime.backends.vllm_utils.deployment import start_rollout_servers +from vime.ray.placement_group import _create_placement_group +from vime.rollout import vllm_rollout +from vime.utils.http_utils import init_http_client, is_port_available, post +from vime.utils.processing_utils import load_tokenizer + +NUM_GPUS = 3 +MODEL_NAME = "Qwen2.5-VL-3B-Instruct" +MODEL_REVISION = "66285546d2b821cf421d4f5eb2576359d3770cd3" +TEST_ROOT = Path(os.environ.get("VIME_TEST_ROOT", "/root")) +MODEL_PATH = TEST_ROOT / "models" / MODEL_NAME +MAX_MODEL_LEN = 4096 +MAX_NEW_TOKENS = 32 +SEED = 2525 + +GROUP_OVERRIDES = { + "gpu_memory_utilization": 0.55, + "max_model_len": MAX_MODEL_LEN, + "max_num_seqs": 4, + "enforce_eager": True, + "generation_config": "vllm", +} + + +def prepare() -> None: + MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) + U.exec_command(f"hf download Qwen/{MODEL_NAME} --revision {MODEL_REVISION} --local-dir {MODEL_PATH}") + + +def _write_config(worker_types: tuple[str, ...]) -> str: + config = { + "vllm": [ + { + "name": "default", + "model_path": str(MODEL_PATH), + "update_weights": False, + "server_groups": [ + { + "worker_type": worker_type, + "num_gpus": 1, + "num_gpus_per_engine": 1, + "overrides": dict(GROUP_OVERRIDES), + } + for worker_type in worker_types + ], + } + ] + } + handle = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", prefix="vime_epd_", delete=False) + with handle: + yaml.safe_dump(config, handle) + return handle.name + + +def _make_args(config_path: str, worker_types: tuple[str, ...]): + args = vllm_parse_args() + args.vllm_config = config_path + args.rollout_external = False + args.rollout_num_gpus = len(worker_types) + args.rollout_num_gpus_per_engine = 1 + args.rollout_num_engines = max(1, sum(worker_type != "encoder" for worker_type in worker_types)) + args.num_gpus_per_node = NUM_GPUS + args.debug_train_only = False + args.debug_rollout_only = True + args.colocate = False + args.actor_num_nodes = 0 + args.actor_num_gpus_per_node = 0 + args.offload_rollout = False + args.use_critic = False + args.critic_num_nodes = 0 + args.critic_num_gpus_per_node = 0 + args.hf_checkpoint = str(MODEL_PATH) + args.seed = SEED + args.fp16 = False + args.use_rollout_routing_replay = False + args.rollout_max_context_len = MAX_MODEL_LEN + args.use_distributed_post = False + args.vllm_server_concurrency = 4 + args.vllm_enable_deterministic_inference = True + args.vllm_router_ip = None + args.vllm_router_port = None + args.vllm_pipeline_parallel_size = 1 + args.vllm_data_parallel_size = 1 + args.vllm_dp_size = 1 + return args + + +def _shutdown_servers(servers: dict[str, Any]) -> None: + engines = [engine for server in servers.values() for engine in server.all_engines if engine is not None] + try: + if engines: + ray.get([engine.shutdown.remote() for engine in engines], timeout=120) + finally: + for engine in engines: + try: + ray.kill(engine, no_restart=True) + except Exception: + pass + + deadline = time.monotonic() + 60 + while not is_port_available(15000) and time.monotonic() < deadline: + time.sleep(1) + + +@contextmanager +def _deployment(pg, worker_types: tuple[str, ...]): + config_path = _write_config(worker_types) + args = _make_args(config_path, worker_types) + servers: dict[str, Any] = {} + ec_path: Path | None = None + try: + servers, init_handles = start_rollout_servers(args, pg) + if init_handles: + ray.get(init_handles) + init_http_client(args) + + for group in servers["default"].server_groups: + config = group.vllm_overrides.get("ec_transfer_config") or {} + extra = config.get("ec_connector_extra_config") or {} + if path := extra.get("shared_storage_path"): + ec_path = Path(path) + break + yield args, servers["default"] + finally: + try: + _shutdown_servers(servers) + finally: + if ec_path is not None: + shutil.rmtree(ec_path, ignore_errors=True) + Path(config_path).unlink(missing_ok=True) + + +def _image_data_url() -> str: + image = Image.new("RGB", (64, 64)) + pixels = image.load() + for y in range(64): + for x in range(64): + pixels[x, y] = (220, 40, 40) if (x // 16 + y // 16) % 2 == 0 else (30, 90, 220) + output = io.BytesIO() + image.save(output, format="PNG") + return "data:image/png;base64," + base64.b64encode(output.getvalue()).decode("ascii") + + +async def _generate(args, messages: list[dict[str, Any]], *, epd_server=None) -> list[int]: + base_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + await vllm_rollout.prime_encoder(args, messages) + if epd_server is not None: + assert getattr(args, "vllm_model_encoder_endpoints", {}), "EPD deployment did not expose an encoder endpoint" + storage_paths = [ + group.vllm_overrides["ec_transfer_config"]["ec_connector_extra_config"]["shared_storage_path"] + for group in epd_server.server_groups + if group.worker_type == "encoder" + ] + assert any( + Path(path).glob("*/encoder_cache.safetensors") for path in storage_paths + ), "encoder did not publish an external EC cache" + render_data = await post( + f"{base_url}/v1/chat/completions/render", + {"model": str(MODEL_PATH), "messages": messages}, + ) + body = vllm_rollout._mm_render_response_to_generate_body(render_data, str(MODEL_PATH)) + body["sampling_params"] = vllm_rollout._build_inference_sampling_params( + { + "max_new_tokens": MAX_NEW_TOKENS, + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "seed": SEED, + "skip_special_tokens": False, + } + ) + output = await post(f"{base_url}/inference/v1/generate", body) + token_ids = output["choices"][0].get("token_ids") or [] + assert token_ids, output + return [int(token_id) for token_id in token_ids] + + +def _find_config(value: Any, key: str) -> dict[str, Any] | None: + if isinstance(value, dict): + candidate = value.get(key) + if isinstance(candidate, dict): + return candidate + for child in value.values(): + if found := _find_config(child, key): + return found + elif isinstance(value, list): + for child in value: + if found := _find_config(child, key): + return found + return None + + +def _engine_info(server) -> dict[str, tuple[str, dict[str, Any]]]: + result = {} + for group in server.server_groups: + endpoint = ray.get(group.engines[0].get_url.remote()) + response = requests.get(f"{endpoint}/server_info?config_format=json", timeout=30) + response.raise_for_status() + result[group.worker_type] = (endpoint, response.json()) + return result + + +def _validate_epd_services(server) -> None: + info = _engine_info(server) + encoder_ec = _find_config(info["encoder"][1], "ec_transfer_config") + prefill_ec = _find_config(info["prefill"][1], "ec_transfer_config") + prefill_kv = _find_config(info["prefill"][1], "kv_transfer_config") + decode_ec = _find_config(info["decode"][1], "ec_transfer_config") + decode_kv = _find_config(info["decode"][1], "kv_transfer_config") + + assert encoder_ec is not None + assert encoder_ec["ec_connector"] == "ECExampleConnector" + assert encoder_ec["ec_role"] == "ec_producer" + assert prefill_ec is not None + assert prefill_ec["ec_connector"] == "ECExampleConnector" + assert prefill_ec["ec_role"] == "ec_consumer" + assert prefill_kv is not None and prefill_kv["kv_role"] == "kv_producer" + assert decode_ec is None or decode_ec.get("ec_role") is None + assert decode_kv is not None and decode_kv["kv_role"] == "kv_consumer" + + response = requests.get(f"http://{server.router_ip}:{server.router_port}/workers", timeout=30) + response.raise_for_status() + workers = response.json()["workers"] + assert {worker["worker_type"] for worker in workers} == {"prefill", "decode"} + assert info["encoder"][0] not in {worker["url"] for worker in workers} + + config = server.server_groups[0].vllm_overrides["ec_transfer_config"] + path = Path(config["ec_connector_extra_config"]["shared_storage_path"]) + assert path.is_dir() + + +def execute() -> None: + ray.init() + pg = _create_placement_group(NUM_GPUS) + image_messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": _image_data_url()}}, + {"type": "text", "text": "Name the two dominant colors. Answer briefly."}, + ], + } + ] + try: + with _deployment(pg, ("regular",)) as (baseline_args, _): + baseline_tokens = asyncio.run(_generate(baseline_args, image_messages)) + + with _deployment(pg, ("encoder", "prefill", "decode")) as (epd_args, epd_server): + epd_tokens = asyncio.run(_generate(epd_args, image_messages, epd_server=epd_server)) + assert epd_tokens == baseline_tokens + tokenizer = load_tokenizer(str(MODEL_PATH), trust_remote_code=True) + assert tokenizer.decode(epd_tokens, skip_special_tokens=False) == tokenizer.decode( + baseline_tokens, skip_special_tokens=False + ) + + _validate_epd_services(epd_server) + finally: + if pg[0] is not None: + ray.util.remove_placement_group(pg[0]) + ray.shutdown() + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen2.5_0.5B_async_short.py b/tests/test_qwen2_5_0_5B_non_colocate_pp.py similarity index 64% rename from tests/test_qwen2.5_0.5B_async_short.py rename to tests/test_qwen2_5_0_5B_non_colocate_pp.py index fda3cdbe5..82ce32366 100644 --- a/tests/test_qwen2.5_0.5B_async_short.py +++ b/tests/test_qwen2_5_0_5B_non_colocate_pp.py @@ -1,6 +1,10 @@ +"""E2E smoke test for non-colocated PP distributed weight updates.""" + import os + import vime.utils.external_utils.command_utils as U + MODEL_NAME = "Qwen2.5-0.5B-Instruct" MODEL_TYPE = "qwen2.5-0.5B" NUM_GPUS = 4 @@ -9,37 +13,43 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/root/models", + ) -def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " +def execute(): rollout_args = ( - "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " - "--input-key prompt " + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " "--label-key label " "--apply-chat-template " "--rollout-shuffle " - "--rm-type deepscaler " - "--num-rollout 3 " + "--rm-type math " + "--num-rollout 2 " "--rollout-batch-size 4 " "--n-samples-per-prompt 4 " - "--rollout-max-response-len 8192 " + "--rollout-max-response-len 512 " "--rollout-temperature 0.8 " + "--over-sampling-batch-size 8 " + "--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " "--global-batch-size 16 " - "--balance-data " ) perf_args = ( "--tensor-model-parallel-size 1 " "--sequence-parallel " - "--pipeline-model-parallel-size 1 " + "--pipeline-model-parallel-size 2 " "--context-parallel-size 1 " "--expert-model-parallel-size 1 " "--expert-tensor-parallel-size 1 " "--use-dynamic-batch-size " - "--max-tokens-per-gpu 9216 " + "--max-tokens-per-gpu 4096 " ) grpo_args = ( @@ -62,20 +72,14 @@ def execute(): ) vllm_args = ( + "--rollout-num-gpus 2 " "--rollout-num-gpus-per-engine 1 " - "--vllm-gpu-memory-utilization 0.65 " + "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-cudagraph-capture-size 16 " ) ci_args = "--ci-test " - fault_tolerance_args = ( - "--use-fault-tolerance " - "--rollout-health-check-interval 5 " - "--rollout-health-check-timeout 10 " - "--rollout-health-check-first-wait 0 " - ) - misc_args = ( "--attention-dropout 0.0 " "--hidden-dropout 0.0 " @@ -83,11 +87,15 @@ def execute(): "--attention-softmax-in-fp32 " "--attention-backend flash " "--actor-num-nodes 1 " - "--actor-num-gpus-per-node 1 " - "--rollout-num-gpus 3 " - "--megatron-to-hf-mode bridge " + "--actor-num-gpus-per-node 2 " ) + torch_dist_checkpoint = f"/root/models/{MODEL_NAME}_torch_dist" + ckpt_args = ( + f"--hf-checkpoint /root/models/{MODEL_NAME}/ " + f"--load {torch_dist_checkpoint} " + f"--ref-load {torch_dist_checkpoint} " + ) train_args = ( f"{ckpt_args} " f"{rollout_args} " @@ -97,7 +105,6 @@ def execute(): f"{perf_args} " f"{vllm_args} " f"{ci_args} " - f"{fault_tolerance_args} " f"{misc_args} " ) @@ -105,14 +112,11 @@ def execute(): train_args=train_args, num_gpus_per_node=NUM_GPUS, megatron_model_type=MODEL_TYPE, - train_script="train_async.py", ) if __name__ == "__main__": prepare() - os.environ.pop("http_proxy", None) - os.environ.pop("https_proxy", None) - os.environ.pop("HTTP_PROXY", None) - os.environ.pop("HTTPS_PROXY", None) + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) execute() diff --git a/tests/test_qwen3.5_0.8B_gsm8k_async_short.py b/tests/test_qwen3.5_0.8B_gsm8k_async_short.py index 8691e6fc5..df952ae59 100644 --- a/tests/test_qwen3.5_0.8B_gsm8k_async_short.py +++ b/tests/test_qwen3.5_0.8B_gsm8k_async_short.py @@ -36,6 +36,8 @@ def execute(): "--n-samples-per-prompt 4 " "--rollout-max-response-len 1024 " "--rollout-temperature 0.8 " + "--rollout-top-k 20 " + "--rollout-top-p 0.95 " "--over-sampling-batch-size 8 " "--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " "--global-batch-size 16 " @@ -57,7 +59,7 @@ def execute(): "--expert-model-parallel-size 1 " "--expert-tensor-parallel-size 1 " "--use-dynamic-batch-size " - "--max-tokens-per-gpu 9216 " + "--max-tokens-per-gpu 2048 " ) grpo_args = ( @@ -80,7 +82,7 @@ def execute(): ) vllm_args = ( - "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.65 " "--vllm-max-cudagraph-capture-size 64" + "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.65 " "--vllm-max-cudagraph-capture-size 32" ) ci_args = "--ci-test " diff --git a/tests/test_qwen3.5_0.8B_gsm8k_short.py b/tests/test_qwen3.5_0.8B_gsm8k_short.py index 2daf829b2..c2410cfad 100644 --- a/tests/test_qwen3.5_0.8B_gsm8k_short.py +++ b/tests/test_qwen3.5_0.8B_gsm8k_short.py @@ -36,6 +36,9 @@ def execute(): "--n-samples-per-prompt 4 " "--rollout-max-response-len 1024 " "--rollout-temperature 0.8 " + "--rollout-top-k 20 " + "--rollout-top-p 0.95 " + "--rollout-data-transport nixl " "--over-sampling-batch-size 8 " "--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " "--global-batch-size 16 " @@ -80,7 +83,7 @@ def execute(): ) vllm_args = ( - "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-cudagraph-capture-size 64" + "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-cudagraph-capture-size 32" ) ci_args = "--ci-test " diff --git a/tests/test_qwen3.5_35B_A3B_npu.py b/tests/test_qwen3.5_35B_A3B_npu.py deleted file mode 100644 index 995c56b56..000000000 --- a/tests/test_qwen3.5_35B_A3B_npu.py +++ /dev/null @@ -1,150 +0,0 @@ -import os -import shlex - -import vime.utils.external_utils.command_utils as U - - -TEST_ROOT = os.environ.get("HF_HOME") or "/root" -MODEL_DIR = f"{TEST_ROOT}/models/Qwen3.5-35B-A3B" -DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k" - - -def prepare(): - models_dir = shlex.quote(f"{TEST_ROOT}/models") - datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets") - model_dir = shlex.quote(MODEL_DIR) - dataset_dir = shlex.quote(DATASET_DIR) - - U.exec_command(f"mkdir -p {models_dir} {datasets_dir}") - U.exec_command(f"hf download Qwen/Qwen3.5-35B-A3B --local-dir {model_dir}") - U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}") - - -def execute(): - model_dir = shlex.quote(MODEL_DIR) - prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") - - # NPU skips torch_dist conversion; HF weights load directly via bridge mode. - checkpoint_args = ( - f"--hf-checkpoint {model_dir} " - f"--load {model_dir} " - f"--ref-load {model_dir} " - "--megatron-to-hf-mode bridge " - "--no-load-optim " - ) - - # Smoke-scaled rollout (num-rollout/batch/n-samples trimmed like test_qwen3_30B_A3B_npu). - rollout_args = ( - f"--prompt-data {prompt_data} " - "--input-key prompt " - "--label-key label " - "--apply-chat-template " - "--rollout-shuffle " - "--rm-type deepscaler " - "--num-rollout 2 " - "--rollout-batch-size 4 " - "--n-samples-per-prompt 4 " - "--rollout-max-response-len 2048 " - "--rollout-temperature 1 " - "--global-batch-size 16 " - "--balance-data " - ) - - # TP=2/EP=8 mirrors scripts/run-qwen3.5-35B-A3B-npu.sh; --qkv-format bshd is - # qwen3.5-specific (not part of MODEL_ARGS, so passed explicitly here). - parallel_args = ( - "--tensor-model-parallel-size 2 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 1 " - "--expert-model-parallel-size 8 " - "--expert-tensor-parallel-size 1 " - "--recompute-granularity full " - "--recompute-method uniform " - "--recompute-num-layers 1 " - "--micro-batch-size 1 " - "--qkv-format bshd " - "--max-tokens-per-gpu 9216 " - ) - - grpo_args = ( - "--advantage-estimator grpo " - "--kl-loss-coef 0.00 " - "--kl-loss-type low_var_kl " - "--kl-coef 0.00 " - "--entropy-coef 0.00 " - "--eps-clip 0.2 " - "--eps-clip-high 0.28 " - ) - - optimizer_args = ( - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - "--optimizer-cpu-offload " - "--overlap-cpu-optimizer-d2h-h2d " - "--use-precision-aware-optimizer " - ) - - vllm_args = ( - "--rollout-num-gpus-per-engine 2 " - "--vllm-gpu-memory-utilization 0.7 " - "--vllm-enable-sleep-mode " - "--vllm-weight-sync-mode native " - "--vllm-enforce-eager " - ) - - model_args = ( - "--attention-dropout 0.0 " - "--hidden-dropout 0.0 " - "--accumulate-allreduce-grads-in-fp32 " - "--attention-softmax-in-fp32 " - "--attention-backend flash " - "--use-flash-attn " - ) - - runtime_args = ( - "--train-backend megatron " - "--actor-num-nodes 1 " - "--actor-num-gpus-per-node 8 " - "--rollout-num-gpus 8 " - "--ci-test " - ) - - train_args = ( - checkpoint_args - + rollout_args - + parallel_args - + grpo_args - + optimizer_args - + vllm_args - + model_args - + runtime_args - ) - # Model architecture (--spec, --attention-output-gate, --moe-shared-expert-gate, - # num-experts, moe-* ...) is injected by sourcing scripts/models/qwen3.5-35B-A3B.sh - # via ${MODEL_ARGS[@]}, so only runtime/training args are passed here. - U.execute_train( - train_args=train_args, - num_gpus_per_node=16, - megatron_model_type="qwen3.5-35B-A3B", - extra_env_vars={ - "DISABLE_L2_CACHE": "1", - "VLLM_USE_AOT_COMPILE": "0", - "ASCEND_CUSTOM_OPP_PATH": "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer", - }, - ) - - -def main(): - prepare() - for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): - os.environ.pop(proxy_var, None) - execute() - - -if __name__ == "__main__": - main() diff --git a/tests/test_qwen3.6_35B_A3B_pd_mooncake.py b/tests/test_qwen3.6_35B_A3B_pd_mooncake.py index d79be87d1..fbb262cac 100644 --- a/tests/test_qwen3.6_35B_A3B_pd_mooncake.py +++ b/tests/test_qwen3.6_35B_A3B_pd_mooncake.py @@ -24,6 +24,7 @@ def prepare(): def execute(): + os.environ.setdefault("VLLM_SSM_CONV_STATE_LAYOUT", "DS") debug_data_path = os.environ.get("DEBUG_ROLLOUT_DATA") or tempfile.mktemp( prefix="qwen3_6_35b_a3b_pd_rollout_", suffix=".pt" ) @@ -98,8 +99,8 @@ def execute(): "--vllm-data-parallel-size 4 " "--vllm-enable-expert-parallel " "--vllm-max-num-seqs 512 " - "--vllm-cudagraph-capture-sizes 1 2 4 8 16 24 32 " - '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":3}\' ' + "--vllm-cudagraph-capture-sizes 5 10 20 40 80 " + '--vllm-speculative-config \'{"method":"mtp","num_speculative_tokens":4}\' ' "--prefill-num-servers 1 " ) diff --git a/tests/test_qwen3_0.6B_parallel_check.py b/tests/test_qwen3_0.6B_parallel_check.py index 4876979b0..2e3c8cfda 100644 --- a/tests/test_qwen3_0.6B_parallel_check.py +++ b/tests/test_qwen3_0.6B_parallel_check.py @@ -9,6 +9,17 @@ MODEL_TYPE = "qwen3-0.6B" NUM_GPUS = 8 +# Cover pure DP scaling, all dense parallel dimensions together, and one +# size-4 case per dimension. +PARALLEL_CONFIGS = ( + # num_gpus, tp, pp, cp + (2, 1, 1, 1), + (8, 2, 2, 2), + (8, 4, 1, 1), + (8, 1, 4, 1), + (8, 1, 1, 4), +) + def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") @@ -61,6 +72,8 @@ def execute(): "--rollout-num-gpus 8 " "--vllm-gpu-memory-utilization 0.8 " "--vllm-max-cudagraph-capture-size 16 " + '--vllm-compilation-config \'{"cudagraph_mode":"FULL_DECODE_ONLY"}\' ' + "--vllm-enable-deterministic-inference " ) ci_args = "--ci-test " @@ -89,42 +102,45 @@ def execute(): f"{misc_args} " ) - for i in range(2): + rollout_data_path = "parallel-check-rollout-data.pt" + for i, calculate_per_token_loss in enumerate((False, True)): + loss_args = "--calculate-per-token-loss " if calculate_per_token_loss else "" + rollout_data_args = ( + f"--save-debug-rollout-data {rollout_data_path} " + if i == 0 + else f"--load-debug-rollout-data {rollout_data_path} " + ) U.execute_train( - train_args=train_args - + ( - f"--save-debug-rollout-data data-{i}.pt " - f"--ci-save-grad-norm grad_norms-{i}.pt " - f"--actor-num-gpus-per-node {NUM_GPUS} " + train_args=( + train_args + + loss_args + + rollout_data_args + + f"--ci-save-grad-norm grad_norms-{i}.pt " + + f"--actor-num-gpus-per-node {NUM_GPUS} " ), num_gpus_per_node=NUM_GPUS, megatron_model_type=MODEL_TYPE, ) - parallel_sizes = [1, 2, 4] - for num_gpus in [8, 4, 2]: - for tp_size in parallel_sizes: - for pp_size in [1, 2, 4]: - for cp_size in parallel_sizes: - if tp_size * pp_size * cp_size > num_gpus: - continue - args = train_args + ( - f"--load-debug-rollout-data data-{i}.pt " - f"--ci-load-grad-norm grad_norms-{i}.pt " - f"--context-parallel-size {cp_size} " - f"--tensor-model-parallel-size {tp_size} " - f"--pipeline-model-parallel-size {pp_size} " - "--sequence-parallel " - f"--actor-num-gpus-per-node {num_gpus} " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 8192 " - ) - - U.execute_train( - train_args=args, - num_gpus_per_node=num_gpus, - megatron_model_type=MODEL_TYPE, - ) - train_args += "--calculate-per-token-loss " + for num_gpus, tp_size, pp_size, cp_size in PARALLEL_CONFIGS: + args = ( + train_args + + loss_args + + f"--load-debug-rollout-data {rollout_data_path} " + + f"--ci-load-grad-norm grad_norms-{i}.pt " + + f"--context-parallel-size {cp_size} " + + f"--tensor-model-parallel-size {tp_size} " + + f"--pipeline-model-parallel-size {pp_size} " + + "--sequence-parallel " + + f"--actor-num-gpus-per-node {num_gpus} " + + "--use-dynamic-batch-size " + + "--max-tokens-per-gpu 8192 " + ) + + U.execute_train( + train_args=args, + num_gpus_per_node=num_gpus, + megatron_model_type=MODEL_TYPE, + ) if __name__ == "__main__": diff --git a/tests/test_qwen3_30B_A3B.py b/tests/test_qwen3_30B_A3B.py index ef3783209..32d9a7ed0 100644 --- a/tests/test_qwen3_30B_A3B.py +++ b/tests/test_qwen3_30B_A3B.py @@ -68,9 +68,6 @@ def execute(): grpo_args = ( "--advantage-estimator gspo " - "--use-kl-loss " - "--kl-loss-coef 0.00 " - "--kl-loss-type low_var_kl " "--kl-coef 0.00 " "--entropy-coef 0.00 " "--eps-clip 4e-4 " @@ -94,7 +91,7 @@ def execute(): "--rollout-num-gpus-per-engine 8 " "--vllm-gpu-memory-utilization 0.8 " "--vllm-max-num-seqs 512 " - "--vllm-max-cudagraph-capture-size 16 " + "--vllm-max-cudagraph-capture-size 32 " ) if USE_DEEPEP: diff --git a/tests/test_qwen3_30B_A3B_npu.py b/tests/test_qwen3_30B_A3B_npu.py index cfe0a932f..e3310096a 100644 --- a/tests/test_qwen3_30B_A3B_npu.py +++ b/tests/test_qwen3_30B_A3B_npu.py @@ -1,5 +1,7 @@ import os import shlex +import sys +import tempfile from pathlib import Path import vime.utils.external_utils.command_utils as U @@ -7,30 +9,34 @@ TEST_ROOT = os.environ.get("HF_HOME") or "/root" MODEL_DIR = f"{TEST_ROOT}/models/Qwen3-30B-A3B" -CHECKPOINT_DIR = f"{TEST_ROOT}/models/Qwen3-30B-A3B_torch_dist" DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k" -def prepare(): +def prepare(torch_dist_ref_load=True): models_dir = shlex.quote(f"{TEST_ROOT}/models") datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets") model_dir = shlex.quote(MODEL_DIR) - checkpoint_dir = shlex.quote(CHECKPOINT_DIR) dataset_dir = shlex.quote(DATASET_DIR) U.exec_command(f"mkdir -p {models_dir} {datasets_dir}") U.exec_command(f"hf download Qwen/Qwen3-30B-A3B --local-dir {model_dir}") U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}") - U.exec_command(f"rm -rf {checkpoint_dir}") + if not torch_dist_ref_load: + return None + + # Retain conversion artifacts for inspection; never delete an existing checkpoint. + checkpoint_path = Path(tempfile.mkdtemp(prefix="Qwen3-30B-A3B_torch_dist_", dir=f"{TEST_ROOT}/models")) + checkpoint_dir = shlex.quote(str(checkpoint_path)) U.exec_command( "source scripts/models/qwen3-30B-A3B.sh && " - "PYTHONPATH=/root/Megatron-LM " - f"torchrun --nproc-per-node 8 tools/convert_hf_to_torch_dist.py " + "TRANSFORMERS_VERBOSITY=error " + f"VIME_PLATFORM=npu PYTHONPATH={shlex.quote(str(U.repo_base_dir))}:/root/Megatron-LM:${{PYTHONPATH:-}} " + f"{shlex.quote(sys.executable)} -m torch.distributed.run --nproc-per-node 8 " + "tools/convert_hf_to_torch_dist.py " "${MODEL_ARGS[@]} " f"--hf-checkpoint {model_dir} --save {checkpoint_dir}" ) - checkpoint_path = Path(CHECKPOINT_DIR) tracker = checkpoint_path / "latest_checkpointed_iteration.txt" assert tracker.read_text().strip() == "release" weight_files = [ @@ -39,14 +45,18 @@ def prepare(): if path.is_file() and path.name != "latest_checkpointed_iteration.txt" ] assert weight_files, f"No checkpoint weights found under {checkpoint_path}" + return str(checkpoint_path) -def execute(): +def execute(torch_dist_checkpoint=None): model_dir = shlex.quote(MODEL_DIR) - checkpoint_dir = shlex.quote(CHECKPOINT_DIR) prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") - checkpoint_args = f"--hf-checkpoint {model_dir} " f"--ref-load {checkpoint_dir} " "--no-load-optim " + checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --ref-load {model_dir} --no-load-optim " + if torch_dist_checkpoint is not None: + checkpoint_args = ( + f"--hf-checkpoint {model_dir} --ref-load {shlex.quote(torch_dist_checkpoint)} --no-load-optim " + ) rollout_args = ( f"--prompt-data {prompt_data} " @@ -103,6 +113,7 @@ def execute(): ) vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " "--rollout-num-gpus-per-engine 4 " "--vllm-enable-sleep-mode " "--vllm-enable-expert-parallel " @@ -150,10 +161,10 @@ def execute(): def main(): - prepare() + checkpoint = prepare(torch_dist_ref_load=os.environ.get("VIME_TEST_TORCH_DIST_REF_LOAD", "1") == "1") for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): os.environ.pop(proxy_var, None) - execute() + execute(checkpoint) if __name__ == "__main__": diff --git a/tests/test_qwen3_30B_A3B_r3.py b/tests/test_qwen3_30B_A3B_r3.py index 38000a0ba..3aa2f7121 100644 --- a/tests/test_qwen3_30B_A3B_r3.py +++ b/tests/test_qwen3_30B_A3B_r3.py @@ -40,6 +40,7 @@ def execute(): "--n-samples-per-prompt 4 " "--rollout-max-response-len 8192 " "--rollout-temperature 1 " + "--rollout-data-transport nixl " "--global-batch-size 16 " "--balance-data " ) @@ -68,9 +69,6 @@ def execute(): grpo_args = ( "--advantage-estimator gspo " - "--use-kl-loss " - "--kl-loss-coef 0.00 " - "--kl-loss-type low_var_kl " "--kl-coef 0.00 " "--entropy-coef 0.00 " "--eps-clip 4e-4 " @@ -94,7 +92,7 @@ def execute(): "--rollout-num-gpus-per-engine 8 " "--vllm-gpu-memory-utilization 0.8 " "--vllm-max-num-seqs 512 " - "--vllm-max-cudagraph-capture-size 16 " + "--vllm-max-cudagraph-capture-size 32 " ) if USE_DEEPEP: diff --git a/tests/test_qwen3_4B_ckpt.py b/tests/test_qwen3_4B_ckpt.py index e038fa1c9..6fa132eb4 100644 --- a/tests/test_qwen3_4B_ckpt.py +++ b/tests/test_qwen3_4B_ckpt.py @@ -5,7 +5,7 @@ import vime.utils.external_utils.command_utils as U -ENABLE_EVAL = bool(int(os.environ.get("SLIME_TEST_ENABLE_EVAL", "1"))) +ENABLE_EVAL = bool(int(os.environ.get("VIME_TEST_ENABLE_EVAL", "1"))) MODEL_NAME = "Qwen3-4B" MODEL_TYPE = "qwen3-4B" diff --git a/tests/test_qwen3_4B_external_pd.py b/tests/test_qwen3_4B_external_pd.py new file mode 100644 index 000000000..6a603dc00 --- /dev/null +++ b/tests/test_qwen3_4B_external_pd.py @@ -0,0 +1,368 @@ +"""E2E test for --rollout-external-engine-addrs with a pure-PD external fleet. + +Spawns two vLLM servers out-of-band on a single GPU box (all tp=1): +- 1 prefill (NIXL ``kv_producer``) +- 1 decode (NIXL ``kv_consumer``) + +and points vime at both via ``--rollout-external-engine-addrs ...``. +The first 4 GPUs train. vime queries ``/server_info`` on each engine to +infer per-engine TP / GPU counts and registers them to its PD-enabled router. + +Weight sync defaults to full disk checkpoints. Set ``VIME_TEST_UPDATE_MODE=delta`` +to publish sparse safetensors, apply them to a host-local checkpoint through +``pull_weights``, and reload that checkpoint without forming an NCCL group with +the trainer. +""" + +import json +import os +import socket +import subprocess +import tempfile +import time +import urllib.request +from pathlib import Path + +import vime.utils.external_utils.command_utils as U + +MODEL_NAME = "Qwen3-4B" +MODEL_TYPE = "qwen3-4B" +TORCH_DIST_CKPT = f"/root/models/{MODEL_NAME}_torch_dist" +NUM_GPUS = 6 +NUM_TRAIN_GPUS = 4 +NUM_PREFILL_ENGINES = 1 +NUM_DECODE_ENGINES = 1 + +EXTERNAL_HOST = "127.0.0.1" +PREFILL_PORTS = [13150] +DECODE_PORTS = [13151] +BOOTSTRAP_PORTS = [13160] +DECODE_BOOTSTRAP_PORTS = [13161] + + +def _get_bond_ipv4(): + net_root = Path("/sys/class/net") + if not net_root.exists(): + return None + + bond_ifaces = [ + path.name for path in net_root.iterdir() if path.name.startswith("bond") and path.name[4:].isdigit() + ] + bond_ifaces.sort(key=lambda name: int(name[4:])) + for iface in bond_ifaces: + try: + output = subprocess.check_output(["ip", "-o", "-4", "addr", "show", "dev", iface], text=True) + except (OSError, subprocess.CalledProcessError): + continue + fields = output.split() + for idx, field in enumerate(fields): + if field == "inet" and idx + 1 < len(fields): + return fields[idx + 1].split("/", 1)[0] + return None + + +def _get_external_host(): + env_value = os.environ.get("VIME_TEST_EXTERNAL_PD_HOST") + if env_value and env_value not in ("127.0.0.1", "localhost"): + return env_value + + bond_host = _get_bond_ipv4() + if bond_host is not None: + return bond_host + + master_addr = os.environ.get("MASTER_ADDR") + if master_addr and master_addr not in ("127.0.0.1", "localhost"): + return master_addr + + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + sock.connect(("8.8.8.8", 80)) + host = sock.getsockname()[0] + if host and not host.startswith("127."): + return host + except OSError: + pass + + return EXTERNAL_HOST + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_TRAIN_GPUS, + dir_dst="/root/models", + ) + + +def _get_gpu_split(): + """Partition visible GPUs: 4 train + 1 prefill + 1 decode.""" + all_gpus = os.environ.get("CUDA_VISIBLE_DEVICES", ",".join(str(i) for i in range(NUM_GPUS))).split(",") + assert len(all_gpus) >= NUM_GPUS, f"Expected at least {NUM_GPUS} GPUs, got {len(all_gpus)}" + train_gpus = all_gpus[:NUM_TRAIN_GPUS] + cursor = NUM_TRAIN_GPUS + prefill_gpus = all_gpus[cursor : cursor + NUM_PREFILL_ENGINES] + cursor += NUM_PREFILL_ENGINES + decode_gpus = all_gpus[cursor : cursor + NUM_DECODE_ENGINES] + return train_gpus, prefill_gpus, decode_gpus + + +def _launch_vllm_server( + *, + gpus: list[str], + port: int, + tp: int, + log_path: str, + disaggregation_mode: str, + disaggregation_bootstrap_port: int | None = None, + external_host: str = EXTERNAL_HOST, +) -> subprocess.Popen: + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus) + env["VLLM_SERVER_DEV_MODE"] = "1" + assert disaggregation_bootstrap_port is not None + env["VLLM_NIXL_SIDE_CHANNEL_HOST"] = external_host + env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(disaggregation_bootstrap_port) + env["UCX_NET_DEVICES"] = os.environ.get("VIME_TEST_DISAGGREGATION_IB_DEVICE", "all").strip() or "all" + + kv_role = { + "prefill": "kv_producer", + "decode": "kv_consumer", + }[disaggregation_mode] + + cmd = [ + "vllm", + "serve", + f"/root/models/{MODEL_NAME}", + "--host", + "0.0.0.0", + "--port", + str(port), + "--served-model-name", + f"/root/models/{MODEL_NAME}", + f"/root/models/{MODEL_NAME}/", + "--tensor-parallel-size", + str(tp), + "--gpu-memory-utilization", + "0.6", + "--trust-remote-code", + "--logprobs-mode", + "processed_logprobs", + "--enable-prompt-tokens-details", + "--enable-server-load-tracking", + "--kv-transfer-config", + json.dumps({"kv_connector": "NixlConnector", "kv_role": kv_role}), + "--weight-transfer-config", + json.dumps({"backend": "nccl"}), + ] + + log_file = open(log_path, "w") + process = subprocess.Popen(cmd, env=env, stdout=log_file, stderr=subprocess.STDOUT) + print( + f"Starting external vllm {disaggregation_mode} server on GPUs {gpus} " + f"port={port} tp={tp} (pid={process.pid}), log: {log_path}" + ) + + # Wait up to ~10 minutes for /server_info to come up. /health + # is unreliable for prefill/decode-only nodes, so we poll /server_info + # — that's what vime's discover_external_engines uses anyway. + deadline = time.time() + 600 + while time.time() < deadline: + if process.poll() is not None: + log_tail = Path(log_path).read_text(errors="replace")[-8000:] + raise RuntimeError( + f"{disaggregation_mode} server exited with code {process.returncode}; {log_path} tail:\n{log_tail}" + ) + try: + req = urllib.request.urlopen(f"http://{external_host}:{port}/server_info?config_format=json", timeout=2) + if req.status == 200: + print(f"External vllm {disaggregation_mode} server is ready on GPUs {gpus}") + return process + except Exception: + pass + time.sleep(5) + + process.kill() + process.wait() + log_tail = Path(log_path).read_text(errors="replace")[-8000:] + raise RuntimeError(f"{disaggregation_mode} server failed to start within timeout; {log_path} tail:\n{log_tail}") + + +def execute(): + update_mode = os.environ.get("VIME_TEST_UPDATE_MODE", "full") + assert update_mode in {"full", "delta"} + train_gpus, prefill_gpus, decode_gpus = _get_gpu_split() + external_host = _get_external_host() + print(f"Using external host for vLLM workers: {external_host}") + processes: list[subprocess.Popen] = [] + + # Restrict CUDA_VISIBLE_DEVICES to training GPUs before Ray starts so + # ray's bundle allocator doesn't try to claim the external vllm GPUs. + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(train_gpus) + + def launch_external_engines(): + for idx, (gpu, port, bootstrap_port) in enumerate( + zip(prefill_gpus, PREFILL_PORTS, BOOTSTRAP_PORTS, strict=True) + ): + processes.append( + _launch_vllm_server( + gpus=[gpu], + port=port, + tp=1, + disaggregation_mode="prefill", + disaggregation_bootstrap_port=bootstrap_port, + external_host=external_host, + log_path=f"/tmp/vllm_external_prefill_{idx}.log", + ) + ) + for idx, (gpu, port, bootstrap_port) in enumerate( + zip(decode_gpus, DECODE_PORTS, DECODE_BOOTSTRAP_PORTS, strict=True) + ): + processes.append( + _launch_vllm_server( + gpus=[gpu], + port=port, + tp=1, + disaggregation_mode="decode", + disaggregation_bootstrap_port=bootstrap_port, + external_host=external_host, + log_path=f"/tmp/vllm_external_decode_{idx}.log", + ) + ) + + disk_dir_cm = tempfile.TemporaryDirectory(prefix=f"vime_external_pd_{update_mode}_disk_") + local_checkpoint_dir_cm = tempfile.TemporaryDirectory(prefix="vime_external_pd_local_checkpoint_") + disk_dir = disk_dir_cm.name + local_checkpoint_dir = local_checkpoint_dir_cm.name + try: + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load {TORCH_DIST_CKPT} " + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 3 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 1024 " + "--rollout-temperature 0.8 " + "--global-batch-size 16 " + ) + + perf_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + # Nonzero entropy coef guarantees a nonzero gradient even when all + # rewards in a group tie (advantages=0), so the delta sync writes + # real sparse files instead of an empty no-op. + "--entropy-coef 0.01 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + # No --rollout-num-gpus / --rollout-num-gpus-per-engine: those are + # inferred from /server_info on each external engine (1 prefill + + # 1 decode, all tp=1). + all_addrs = [f"{external_host}:{port}" for port in (*PREFILL_PORTS, *DECODE_PORTS)] + external_args = "--rollout-external-engine-addrs " + " ".join(all_addrs) + " " + + disk_update_args = ( + f"--update-weight-mode {update_mode} " + "--update-weight-transport disk " + f"--update-weight-disk-dir {disk_dir} " + "--update-weight-disk-keep-files " + ) + if update_mode == "delta": + disk_update_args += ( + "--update-weight-delta-encoding xor " f"--update-weight-local-checkpoint-dir {local_checkpoint_dir} " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {NUM_TRAIN_GPUS} " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{external_args} " + f"{disk_update_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_TRAIN_GPUS, + megatron_model_type=MODEL_TYPE, + before_ray_job_submit=launch_external_engines, + extra_env_vars={ + "no_proxy": f"127.0.0.1,localhost,{external_host}", + "NO_PROXY": f"127.0.0.1,localhost,{external_host}", + }, + ) + + checkpoint_dirs = sorted(Path(disk_dir).glob("weight_v*")) + assert checkpoint_dirs, f"No disk checkpoint directories were written under {disk_dir}" + assert any((path / "model.safetensors.index.json").exists() for path in checkpoint_dirs) + assert any(list(path.glob("*.safetensors")) for path in checkpoint_dirs) + if update_mode == "delta": + indexes = [json.loads((path / "model.safetensors.index.json").read_text()) for path in checkpoint_dirs] + assert all("delta_encoding" in index["metadata"] for index in indexes) + finally: + for p in processes: + if p.poll() is None: + p.kill() + p.wait() + U.exec_command("pkill -9 vllm; true") + disk_dir_cm.cleanup() + local_checkpoint_dir_cm.cleanup() + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen3_4B_npu.py b/tests/test_qwen3_4B_npu.py index 1b7d7cf3b..ea94cf6d1 100644 --- a/tests/test_qwen3_4B_npu.py +++ b/tests/test_qwen3_4B_npu.py @@ -24,13 +24,7 @@ def execute(): model_dir = shlex.quote(MODEL_DIR) prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") - checkpoint_args = ( - f"--hf-checkpoint {model_dir} " - f"--ref-load {model_dir} " - f"--load {model_dir} " - "--no-load-optim " - "--megatron-to-hf-mode bridge " - ) + checkpoint_args = f"--hf-checkpoint {model_dir} --ref-load {model_dir} --load {model_dir} --no-load-optim " rollout_args = ( f"--prompt-data {prompt_data} " @@ -84,8 +78,8 @@ def execute(): ) vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " "--rollout-num-gpus-per-engine 4 " - "--vllm-weight-sync-mode native " "--vllm-enable-sleep-mode " "--vllm-gpu-memory-utilization 0.6 " "--vllm-max-model-len 4096 " diff --git a/tests/test_qwen3_4B_ppo_disaggregate.py b/tests/test_qwen3_4B_ppo_disaggregate.py index 6c736c1e9..ae4b46bea 100644 --- a/tests/test_qwen3_4B_ppo_disaggregate.py +++ b/tests/test_qwen3_4B_ppo_disaggregate.py @@ -51,6 +51,7 @@ def execute(): "--n-samples-per-prompt 4 " "--rollout-max-response-len 8192 " "--rollout-temperature 0.8 " + "--rollout-data-transport nixl " "--global-batch-size 16 " "--balance-data " ) diff --git a/tests/test_qwen3_4B_ppo_train_critic_only.py b/tests/test_qwen3_4B_ppo_train_critic_only.py index dd22d54f2..81e4737e4 100644 --- a/tests/test_qwen3_4B_ppo_train_critic_only.py +++ b/tests/test_qwen3_4B_ppo_train_critic_only.py @@ -101,6 +101,7 @@ def execute(): "--rollout-num-gpus 8 " "--vllm-gpu-memory-utilization 0.8 " "--vllm-max-num-seqs 512 " + "--vllm-max-cudagraph-capture-size 16 " ) ci_args = "--ci-test " diff --git a/tests/test_qwen3_4B_streaming_partial_rollout.py b/tests/test_qwen3_4B_streaming_partial_rollout.py index 3f326d3a5..7543f2f46 100644 --- a/tests/test_qwen3_4B_streaming_partial_rollout.py +++ b/tests/test_qwen3_4B_streaming_partial_rollout.py @@ -94,7 +94,7 @@ def execute(): "--rollout-num-gpus 8 " "--vllm-gpu-memory-utilization 0.8 " "--vllm-max-num-seqs 512 " - "--vllm-max-cudagraph-capture-size 16 " + "--vllm-max-cudagraph-capture-size 32 " ) ci_args = "--ci-test " diff --git a/tests/test_qwen2.5_0.5B_short.py b/tests/test_qwen3_5_0_8B_top_p_cp2.py similarity index 56% rename from tests/test_qwen2.5_0.5B_short.py rename to tests/test_qwen3_5_0_8B_top_p_cp2.py index 1d795aa8a..a880f5fe7 100644 --- a/tests/test_qwen2.5_0.5B_short.py +++ b/tests/test_qwen3_5_0_8B_top_p_cp2.py @@ -1,45 +1,60 @@ +"""Four-GPU Qwen3.5 top-p replay E2E with context parallelism.""" + import os + import vime.utils.external_utils.command_utils as U -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" + +os.environ.setdefault("NCCL_NVLS_ENABLE", "0") + + +MODEL_NAME = "Qwen3.5-0.8B" +MODEL_TYPE = "qwen3.5-0.8B" NUM_GPUS = 4 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} " f"--ref-load {TORCH_DIST_CKPT} " rollout_args = ( - "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " - "--input-key prompt " + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " "--label-key label " "--apply-chat-template " "--rollout-shuffle " - "--rm-type deepscaler " - "--num-rollout 3 " - "--rollout-batch-size 4 " + "--rm-type math " + "--num-rollout 1 " + "--rollout-batch-size 2 " "--n-samples-per-prompt 4 " - "--rollout-max-response-len 8192 " + "--rollout-max-response-len 512 " "--rollout-temperature 0.8 " - "--global-batch-size 16 " - "--balance-data " + "--rollout-top-k 20 " + "--rollout-top-p 0.95 " + "--global-batch-size 8 " ) perf_args = ( "--tensor-model-parallel-size 1 " "--sequence-parallel " "--pipeline-model-parallel-size 1 " - "--context-parallel-size 1 " + "--context-parallel-size 2 " "--expert-model-parallel-size 1 " "--expert-tensor-parallel-size 1 " "--use-dynamic-batch-size " - "--max-tokens-per-gpu 9216 " + "--max-tokens-per-gpu 4096 " ) grpo_args = ( @@ -65,42 +80,29 @@ def execute(): "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-cudagraph-capture-size 16 " ) - ci_args = "--ci-test " - - fault_tolerance_args = ( - "--use-fault-tolerance " - "--rollout-health-check-interval 5 " - "--rollout-health-check-timeout 10 " - "--rollout-health-check-first-wait 0 " - ) - misc_args = ( "--attention-dropout 0.0 " "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " + "--loss-mask-type qwen3_5 " "--actor-num-nodes 1 " "--actor-num-gpus-per-node 4 " "--colocate " - "--megatron-to-hf-mode bridge " - ) - - train_args = ( - f"{ckpt_args} " - f"{rollout_args} " - f"{optimizer_args} " - f"{grpo_args} " - f"{U.get_default_wandb_args(__file__)} " - f"{perf_args} " - f"{vllm_args} " - f"{ci_args} " - f"{fault_tolerance_args} " - f"{misc_args} " + "--ci-test " ) U.execute_train( - train_args=train_args, + train_args=( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{perf_args} " + f"{vllm_args} " + f"{misc_args} " + ), num_gpus_per_node=NUM_GPUS, megatron_model_type=MODEL_TYPE, ) @@ -108,8 +110,6 @@ def execute(): if __name__ == "__main__": prepare() - os.environ.pop("http_proxy", None) - os.environ.pop("https_proxy", None) - os.environ.pop("HTTP_PROXY", None) - os.environ.pop("HTTPS_PROXY", None) + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) execute() diff --git a/tests/test_qwen3_5_mtp_bridge_mapping.py b/tests/test_qwen3_5_mtp_bridge_mapping.py deleted file mode 100644 index d1267f649..000000000 --- a/tests/test_qwen3_5_mtp_bridge_mapping.py +++ /dev/null @@ -1,242 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest -import torch - - -def install_bridge_stubs(): - megatron_mod = types.ModuleType("megatron") - core_mod = types.ModuleType("megatron.core") - models_mod = types.ModuleType("megatron.core.models") - gpt_mod = types.ModuleType("megatron.core.models.gpt") - gpt_layer_specs_mod = types.ModuleType("megatron.core.models.gpt.gpt_layer_specs") - gpt_layer_specs_mod.get_gpt_mtp_block_spec = lambda _config, transformer_layer_spec, **_kwargs: ( - "mtp-spec", - transformer_layer_spec, - ) - - mbridge_mod = types.ModuleType("mbridge") - mbridge_core_mod = types.ModuleType("mbridge.core") - mbridge_models_mod = types.ModuleType("mbridge.models") - - def register_model(_names): - def decorator(cls): - return cls - - return decorator - - class Qwen2MoEBridge: - _MLP_MAPPING = { - "shared_experts.linear_fc1.weight": [ - "model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight", - "model.layers.{layer_number}.mlp.shared_expert.up_proj.weight", - ], - "pre_mlp_layernorm": ["model.layers.{layer_number}.post_attention_layernorm.weight"], - "shared_experts.linear_fc2.weight": ["model.layers.{layer_number}.mlp.shared_expert.down_proj.weight"], - "mlp.router.weight": ["model.layers.{layer_number}.mlp.gate.weight"], - "shared_experts.gate_weight": ["model.layers.{layer_number}.mlp.shared_expert_gate.weight"], - "mlp.experts.linear_fc1": [ - "model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", - "model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", - ], - "mlp.experts.linear_fc2": ["model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight"], - } - - def _weight_name_mapping_mlp(self, name: str) -> list[str]: - layer_number = name.split(".")[2] - convert_names = [] - for keyword, mapping_names in self._MLP_MAPPING.items(): - if keyword in name: - if "{expert_id}" in mapping_names[0]: - expert_id = name.split("weight")[-1] - convert_names.extend( - [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names] - ) - else: - convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names]) - break - if len(convert_names) == 0: - raise NotImplementedError(f"Unsupported parameter name: {name}") - return convert_names - - def _weight_name_mapping_attention(self, name: str) -> list[str]: - raise NotImplementedError(f"Unexpected attention mapping lookup: {name}") - - def _get_transformer_layer_spec(self, vp_stage=None): - return "REAL_LAYER_SPEC" if vp_stage is None else f"REAL_LAYER_SPEC_VP{vp_stage}" - - def _get_gptmodel_args(self) -> dict: - return {"base": "ok"} - - def _model_provider(self, callbacks): - def provider(pre_process, post_process, vp_stage=None): - transformer_layer_spec = self._get_transformer_layer_spec(vp_stage) - gptmodel_args = self._get_gptmodel_args() - return {"transformer_layer_spec": transformer_layer_spec, **gptmodel_args} - - return provider - - def _weight_to_mcore_format(self, _mcore_weights_name, hf_weights): - assert len(hf_weights) == 1 - return hf_weights[0] - - def _weight_to_hf_format(self, mcore_weights_name, mcore_weights): - return [mcore_weights_name], [mcore_weights] - - def _build_base_config(self, **kwargs): - return kwargs - - mbridge_core_mod.register_model = register_model - mbridge_models_mod.Qwen2MoEBridge = Qwen2MoEBridge - - sys.modules["megatron"] = megatron_mod - sys.modules["megatron.core"] = core_mod - sys.modules["megatron.core.models"] = models_mod - sys.modules["megatron.core.models.gpt"] = gpt_mod - sys.modules["megatron.core.models.gpt.gpt_layer_specs"] = gpt_layer_specs_mod - sys.modules["mbridge"] = mbridge_mod - sys.modules["mbridge.core"] = mbridge_core_mod - sys.modules["mbridge.models"] = mbridge_models_mod - - -def load_bridge_module(): - install_bridge_stubs() - module_path = Path(__file__).resolve().parents[1] / "vime_plugins" / "mbridge" / "qwen3_5.py" - module_name = "test_qwen3_5_bridge_module" - sys.modules.pop(module_name, None) - spec = importlib.util.spec_from_file_location(module_name, module_path) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -def load_raw_export_module(): - module_path = ( - Path(__file__).resolve().parents[1] / "vime" / "backends" / "megatron_utils" / "megatron_to_hf" / "qwen3_5.py" - ) - module_name = "test_qwen3_5_raw_export_module" - sys.modules.pop(module_name, None) - spec = importlib.util.spec_from_file_location(module_name, module_path) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -@pytest.mark.unit -def test_mtp_moe_expert_mapping_uses_individual_hf_weights(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - - fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight42") - fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight42") - - assert fc1_names == [ - "mtp.layers.0.mlp.experts.42.gate_proj.weight", - "mtp.layers.0.mlp.experts.42.up_proj.weight", - ] - assert fc2_names == ["mtp.layers.0.mlp.experts.42.down_proj.weight"] - - -@pytest.mark.unit -def test_mtp_dense_mlp_mapping_still_uses_dense_hf_weights(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - - fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc1.weight") - fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc2.weight") - - assert fc1_names == ["mtp.layers.0.mlp.gate_proj.weight", "mtp.layers.0.mlp.up_proj.weight"] - assert fc2_names == ["mtp.layers.0.mlp.down_proj.weight"] - - -@pytest.mark.unit -def test_mtp_block_spec_uses_current_transformer_layer_spec(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - bridge.config = "CONFIG_OBJECT" - bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1)) - - provider = bridge._model_provider([]) - result = provider(True, True, vp_stage=3) - - assert result["transformer_layer_spec"] == "REAL_LAYER_SPEC_VP3" - assert result["mtp_block_spec"] == ("mtp-spec", "REAL_LAYER_SPEC_VP3") - - -@pytest.mark.unit -def test_tied_qwen3_5_uses_language_embedding_for_output_layer(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(tie_word_embeddings=True)) - - bridge._adjust_mapping_for_shared_weights() - - assert bridge._DIRECT_MAPPING["output_layer.weight"] == "model.language_model.embed_tokens.weight" - assert module.Qwen3_5Bridge._DIRECT_MAPPING["output_layer.weight"] == "lm_head.weight" - - -@pytest.mark.unit -def test_eh_proj_keeps_column_order_when_loading_to_mcore(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - - weight = torch.arange(24, dtype=torch.float32).view(3, 8) - converted = bridge._weight_to_mcore_format("mtp.layers.0.eh_proj.weight", [weight]) - - assert torch.equal(converted, weight) - - -@pytest.mark.unit -def test_build_config_enables_gated_attention_when_transformer_config_supports_it(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1)) - bridge.TransformerConfigClass = types.SimpleNamespace( - __dataclass_fields__={ - "mtp_num_layers": None, - "attention_output_gate": None, - "use_gated_attention": None, - } - ) - - config = bridge._build_config() - - assert config["mtp_num_layers"] == 1 - assert config["attention_output_gate"] is True - assert config["use_gated_attention"] is True - - -@pytest.mark.unit -def test_build_config_skips_gated_attention_when_transformer_config_does_not_support_it(): - module = load_bridge_module() - bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1)) - bridge.TransformerConfigClass = types.SimpleNamespace( - __dataclass_fields__={ - "mtp_num_layers": None, - "attention_output_gate": None, - } - ) - - config = bridge._build_config() - - assert config["mtp_num_layers"] == 1 - assert config["attention_output_gate"] is True - assert "use_gated_attention" not in config - - -@pytest.mark.unit -def test_raw_qwen3_5_mtp_export_keeps_eh_proj_column_order(): - module = load_raw_export_module() - - weight = torch.arange(24, dtype=torch.float32).view(3, 8) - converted = module.convert_qwen3_5_to_hf( - types.SimpleNamespace(), "module.module.mtp.layers.0.eh_proj.weight", weight - ) - - assert converted == [("mtp.fc.weight", weight)] diff --git a/tests/test_qwen3_5_vl_native.py b/tests/test_qwen3_5_vl_native.py new file mode 100644 index 000000000..aac94b833 --- /dev/null +++ b/tests/test_qwen3_5_vl_native.py @@ -0,0 +1,206 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +# Keep this CPU test independent of the Megatron runtime package. The tested +# mapping and packed-sequence helpers themselves only depend on torch. +try: + _has_megatron = importlib.util.find_spec("megatron.core") is not None +except ModuleNotFoundError: + _has_megatron = False +if not _has_megatron: + _megatron_utils = types.ModuleType("vime.backends.megatron_utils") + _megatron_utils.__path__ = [str(Path(__file__).resolve().parents[1] / "vime/backends/megatron_utils")] + sys.modules["vime.backends.megatron_utils"] = _megatron_utils + +from vime.backends.megatron_utils.hf_to_megatron.common import SafetensorReader, _tensor_parallel_shard +from vime.backends.megatron_utils.hf_to_megatron.qwen3_5 import qwen3_5_hf_tensor +from vime.backends.megatron_utils.megatron_to_hf.qwen3_5 import convert_qwen3_5_to_hf +from vime_plugins.models.qwen3_5_vl_utils import build_packed_mrope_position_ids, get_packed_cp_local_indices + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_packed_mrope_resets_positions_for_each_sample(): + input_ids = torch.tensor( + [[99, 10, 10, 10, 10, 98, 7, 99, 10, 10, 10, 10, 98, 8]], + dtype=torch.long, + ) + positions = build_packed_mrope_position_ids( + input_ids, + [0, 7, 14], + image_grid_thw=torch.tensor([[1, 4, 4], [1, 4, 4]]), + video_grid_thw=None, + image_token_id=10, + video_token_id=20, + vision_start_token_id=99, + spatial_merge_size=2, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4], + [0, 1, 1, 2, 2, 3, 4], + [0, 1, 2, 1, 2, 3, 4], + ] + ) + assert torch.equal(positions[:, 0, :7], expected) + assert torch.equal(positions[:, 0, 7:], expected) + + +@pytest.mark.unit +def test_packed_mrope_rejects_unused_grids(): + with pytest.raises(ValueError, match="Unused .* image grids"): + build_packed_mrope_position_ids( + torch.tensor([[1, 2, 3]]), + [0, 3], + image_grid_thw=torch.tensor([[1, 4, 4]]), + video_grid_thw=None, + image_token_id=10, + video_token_id=20, + vision_start_token_id=99, + spatial_merge_size=2, + ) + + +@pytest.mark.unit +def test_thd_cp_indices_select_two_chunks_per_packed_sequence(): + rank_0 = get_packed_cp_local_indices([0, 8, 16], cp_size=2, cp_rank=0, device=torch.device("cpu")) + rank_1 = get_packed_cp_local_indices([0, 8, 16], cp_size=2, cp_rank=1, device=torch.device("cpu")) + + assert rank_0.tolist() == [0, 1, 6, 7, 8, 9, 14, 15] + assert rank_1.tolist() == [2, 3, 4, 5, 10, 11, 12, 13] + + +@pytest.mark.unit +def test_raw_qkv_loader_is_inverse_of_exporter(): + hidden_size = 3 + q = torch.arange(16 * hidden_size).reshape(16, hidden_size) + k = torch.arange(4 * hidden_size).reshape(4, hidden_size) + 1000 + v = torch.arange(4 * hidden_size).reshape(4, hidden_size) + 2000 + tensors = { + "model.language_model.layers.0.self_attn.q_proj.weight": q, + "model.language_model.layers.0.self_attn.k_proj.weight": k, + "model.language_model.layers.0.self_attn.v_proj.weight": v, + } + reader = types.SimpleNamespace(get_tensor=tensors.__getitem__) + text_config = types.SimpleNamespace( + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + ) + hf_config = types.SimpleNamespace(text_config=text_config, tie_word_embeddings=False) + + mcore_qkv = qwen3_5_hf_tensor( + "module.module.language_model.decoder.layers.0.self_attention.linear_qkv.weight", + reader, + hf_config, + ) + converted = dict( + convert_qwen3_5_to_hf( + types.SimpleNamespace( + kv_channels=2, + hidden_size=hidden_size, + num_attention_heads=4, + num_query_groups=2, + ), + "module.module.language_model.decoder.layers.0.self_attention.linear_qkv.weight", + mcore_qkv, + ) + ) + + assert torch.equal(converted["model.language_model.layers.0.self_attn.q_proj.weight"], q) + assert torch.equal(converted["model.language_model.layers.0.self_attn.k_proj.weight"], k) + assert torch.equal(converted["model.language_model.layers.0.self_attn.v_proj.weight"], v) + + +@pytest.mark.unit +def test_raw_loader_uses_hf_vision_name_directly(): + weight = torch.randn(4, 3) + key = "model.visual.patch_embed.proj.weight" + reader = types.SimpleNamespace(get_tensor={key: weight}.__getitem__) + hf_config = types.SimpleNamespace(tie_word_embeddings=False) + + loaded = qwen3_5_hf_tensor( + "module.module.model.visual.patch_embed.proj.weight", + reader, + hf_config, + ) + assert loaded is weight + + +@pytest.mark.unit +def test_raw_loader_uses_individual_mtp_expert_weights(): + gate = torch.randn(4, 3) + up = torch.randn(4, 3) + down = torch.randn(3, 4) + tensors = { + "mtp.layers.0.mlp.experts.7.gate_proj.weight": gate, + "mtp.layers.0.mlp.experts.7.up_proj.weight": up, + "mtp.layers.0.mlp.experts.7.down_proj.weight": down, + } + reader = types.SimpleNamespace(get_tensor=tensors.__getitem__) + config = types.SimpleNamespace( + text_config=types.SimpleNamespace(tie_word_embeddings=False), + tie_word_embeddings=False, + ) + + fc1 = qwen3_5_hf_tensor( + "module.module.language_model.mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight7", + reader, + config, + ) + fc2 = qwen3_5_hf_tensor( + "module.module.language_model.mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight7", + reader, + config, + ) + + assert torch.equal(fc1, torch.cat((gate, up))) + assert fc2 is down + + +@pytest.mark.unit +def test_raw_loader_shards_swiglu_and_grouped_moe_fc2(): + fc1 = torch.arange(8 * 3).reshape(8, 3) + fc1_shard = _tensor_parallel_shard( + "module.module.language_model.decoder.layers.0.mlp.linear_fc1.weight", + fc1, + parallel_size=2, + parallel_rank=1, + partition_dim=0, + partition_stride=1, + ) + assert torch.equal(fc1_shard, torch.cat((fc1[2:4], fc1[6:8]))) + + fc2 = torch.arange(4 * 6).reshape(4, 6) + fc2_shard = _tensor_parallel_shard( + "module.module.language_model.decoder.layers.0.mlp.experts.linear_fc2.weight0", + fc2, + parallel_size=2, + parallel_rank=1, + partition_dim=0, + partition_stride=1, + ) + assert torch.equal(fc2_shard, fc2[:, 3:]) + + +@pytest.mark.unit +def test_safetensor_reader_caches_only_the_last_tensor(tmp_path): + save_file({"first": torch.ones(2), "second": torch.zeros(2)}, tmp_path / "model.safetensors") + reader = SafetensorReader(tmp_path) + + first = reader.get_tensor("first") + assert reader.get_tensor("first") is first + reader.get_tensor("second") + assert reader.get_tensor("first") is not first + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_qwen3_5_vl_train_rollout_e2e.py b/tests/test_qwen3_5_vl_train_rollout_e2e.py new file mode 100644 index 000000000..1cc311dfc --- /dev/null +++ b/tests/test_qwen3_5_vl_train_rollout_e2e.py @@ -0,0 +1,125 @@ +"""Eight-GPU Qwen3.5-VL rollout, train, and weight-update E2E.""" + +import os + +import vime.utils.external_utils.command_utils as U + + +os.environ.setdefault("NCCL_NVLS_ENABLE", "0") + + +MODEL_NAME = "Qwen3.5-35B-A3B" +MODEL_TYPE = "qwen3.5-35B-A3B-vl" +NUM_GPUS = 8 +DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed" +DATASET_ROOT = "/root/datasets/geo3k_imgurl_processed" +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download --repo-type dataset {DATASET_NAME} --local-dir {DATASET_ROOT}") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} " f"--ref-load {TORCH_DIST_CKPT} " + + rollout_args = ( + f"--prompt-data {DATASET_ROOT}/train.parquet " + "--input-key problem " + "--label-key answer " + '--multimodal-keys \'{"image": "images"}\' ' + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 1 " + "--rollout-batch-size 2 " + "--n-samples-per-prompt 2 " + "--rollout-max-response-len 512 " + "--rollout-temperature 0.8 " + "--global-batch-size 4 " + ) + + perf_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--micro-batch-size 1 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 8 " + "--vllm-enable-expert-parallel " + "--vllm-gpu-memory-utilization 0.6 " + "--vllm-max-model-len 4096 " + "--vllm-max-num-seqs 4 " + "--vllm-enforce-eager " + "--vllm-generation-config vllm " + "--vllm-logprobs-mode processed_logprobs " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--loss-mask-type qwen3_5 " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--ci-test " + ) + + U.execute_train( + train_args=( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{perf_args} " + f"{vllm_args} " + f"{misc_args} " + ), + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen3_linear_attention_cu_seqlens.py b/tests/test_qwen3_linear_attention_cu_seqlens.py index 6abb3cc9d..75675198d 100644 --- a/tests/test_qwen3_linear_attention_cu_seqlens.py +++ b/tests/test_qwen3_linear_attention_cu_seqlens.py @@ -9,6 +9,8 @@ import torch import torch.nn as nn +NUM_GPUS = 0 + def install_megatron_stubs() -> None: if "megatron" in sys.modules: @@ -146,7 +148,7 @@ def test_linear_attention_forwards_cu_seqlens_to_chunk_kernel( ): module = load_module(module_name) - monkeypatch.setattr(module.torch.cuda, "current_device", lambda: "cpu") + monkeypatch.setattr(module.accelerator, "current_device", lambda: "cpu") monkeypatch.setattr(module, "ShortConvolution", FakeShortConvolution, raising=False) monkeypatch.setattr(module, "FusedRMSNormGated", FakeFusedRMSNormGated, raising=False) @@ -188,3 +190,7 @@ def fake_get_chunk_gated_delta_rule(backend): assert output.shape == hidden_states.shape assert len(chunk_calls) == 1 assert torch.equal(chunk_calls[0], cu_seqlens) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_qwen3_vl_8B_npu.py b/tests/test_qwen3_vl_8B_npu.py index ff0577753..c117a6744 100644 --- a/tests/test_qwen3_vl_8B_npu.py +++ b/tests/test_qwen3_vl_8B_npu.py @@ -5,7 +5,7 @@ # Single-turn Qwen3-VL GRPO on geo3k (mirrors examples/geo3k_vlm/run_geo3k_vlm_npu.sh). -# Qwen3-VL-8B maps to the qwen3-8B megatron config; the vision tower is handled by bridge. +# Qwen3-VL-8B uses the qwen3-8B language config and the native VL provider. MODEL_NAME = "Qwen3-VL-8B-Instruct" MODEL_TYPE = "qwen3-8B" TEST_ROOT = os.environ.get("HF_HOME") or "/root" @@ -28,9 +28,7 @@ def execute(): model_dir = shlex.quote(MODEL_DIR) prompt_data = shlex.quote(f"{DATASET_DIR}/train.parquet") - checkpoint_args = ( - f"--hf-checkpoint {model_dir} " f"--load {model_dir} " "--megatron-to-hf-mode bridge " "--no-load-optim " - ) + checkpoint_args = f"--hf-checkpoint {model_dir} --load {model_dir} --no-load-optim " rollout_args = ( f"--prompt-data {prompt_data} " @@ -82,14 +80,16 @@ def execute(): ) vllm_args = ( + "--vllm-additional-config '{\"weight_nz_mode\":0}' " "--rollout-num-gpus-per-engine 1 " - "--vllm-gpu-memory-utilization 0.8 " + "--vllm-gpu-memory-utilization 0.7 " "--vllm-max-model-len 16384 " "--vllm-generation-config auto " "--vllm-logprobs-mode processed_logprobs " ) model_args = ( + "--spec vime_plugins.models.qwen3_vl get_qwen3_vl_model_provider " "--attention-dropout 0.0 " "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " diff --git a/tests/test_qwen3_vl_native.py b/tests/test_qwen3_vl_native.py new file mode 100644 index 000000000..135fcd5b5 --- /dev/null +++ b/tests/test_qwen3_vl_native.py @@ -0,0 +1,318 @@ +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +NUM_GPUS = 0 +pytestmark = pytest.mark.unit + + +@pytest.fixture +def native(monkeypatch): + monkeypatch.setenv("VIME_PLATFORM", "cuda") + pytest.importorskip("megatron.core") + from vime_plugins.models import qwen3_vl + + return qwen3_vl + + +@pytest.fixture +def hf_config(): + from transformers import Qwen3VLConfig + + return Qwen3VLConfig( + image_token_id=10, + video_token_id=20, + vision_start_token_id=30, + text_config={ + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "vocab_size": 32, + "rope_scaling": {"rope_type": "default", "mrope_section": [1, 1, 0], "mrope_interleaved": True}, + "rope_theta": 5000000, + }, + vision_config={ + "hidden_size": 8, + "intermediate_size": 16, + "depth": 2, + "num_heads": 2, + "out_hidden_size": 8, + "patch_size": 2, + "temporal_patch_size": 2, + "spatial_merge_size": 2, + "num_position_embeddings": 4, + "deepstack_visual_indexes": [0], + }, + ) + + +def test_vision_native_load_export_and_backward(native, hf_config, tmp_path, monkeypatch): + from vime.backends.megatron_utils.hf_to_megatron.common import load_model_hf_weights + from vime.backends.megatron_utils.hf_to_megatron.qwen3_vl import qwen3_vl_hf_tensor + from vime.backends.megatron_utils.megatron_to_hf import convert_to_hf + from vime.backends.megatron_utils.update_weight import common + + config = SimpleNamespace(use_cpu_initialization=True, params_dtype=torch.float32, recompute_granularity="full") + vision = native._load_vision_model(hf_config, config) + tensors = {"model.visual." + name: param.detach().clone() for name, param in vision.named_parameters()} + save_file(tensors, tmp_path / "model.safetensors") + with torch.no_grad(): + for parameter in vision.parameters(): + parameter.zero_() + named = [("module.module.model.visual." + name, param) for name, param in vision.named_parameters()] + monkeypatch.setattr(common, "named_params_and_buffers", lambda args, model: iter(named)) + load_model_hf_weights(SimpleNamespace(), [vision], tmp_path, hf_config, qwen3_vl_hf_tensor) + + for name, parameter in named: + [(hf_name, exported)] = convert_to_hf(None, "qwen3_vl", name, parameter) + assert torch.equal(exported, tensors[hf_name]) + assert parameter.requires_grad + assert not parameter.tensor_model_parallel + assert parameter.partition_dim == -1 + output = vision(torch.randn(4, 24), grid_thw=torch.tensor([[1, 2, 2]])) + features, deepstack = ( + (output.pooler_output, output.deepstack_features) if hasattr(output, "pooler_output") else output + ) + (features.square().sum() + deepstack[0].square().sum()).backward() + assert vision.patch_embed.proj.weight.grad.abs().sum() > 0 + assert vision.deepstack_merger_list[0].linear_fc2.weight.grad.abs().sum() > 0 + + +def _injection_model(native, *, sequence_parallel=False): + class Embedding: + def __call__(self, input_ids, position_ids): + return input_ids.T[..., None].float().expand(-1, -1, 2).clone() + + class Vision: + dtype = torch.float32 + + def __call__(self, values, grid_thw): + return SimpleNamespace(pooler_output=values, deepstack_features=[values * 2, values * 3]) + + return SimpleNamespace( + config=SimpleNamespace(sequence_parallel=sequence_parallel), + language_model=SimpleNamespace(embedding=Embedding()), + model=SimpleNamespace(visual=Vision()), + image_token_id=10, + video_token_id=20, + ) + + +def test_vision_and_deepstack_follow_token_order_and_keep_gradients(native): + model = _injection_model(native) + image = torch.tensor([[100.0, 101.0]], requires_grad=True) + video = torch.tensor([[200.0, 201.0]], requires_grad=True) + embeddings, mask, deepstack = native.Qwen3VLModel._inject_vision_embeddings( + model, torch.tensor([[7, 20, 8, 10]]), image, video, torch.tensor([[1, 2, 2]]), torch.tensor([[1, 2, 2]]) + ) + expected = torch.cat((video, image)) + assert torch.equal(embeddings[:, 0][mask[0]], expected) + assert torch.equal(deepstack[0], expected * 2) + (embeddings.sum() + sum(t.sum() for t in deepstack)).backward() + assert torch.equal(image.grad, torch.full_like(image, 6)) + assert torch.equal(video.grad, torch.full_like(video, 6)) + + +def test_deepstack_sp_routes_gradients_through_tp_copy(native, monkeypatch): + copied = [] + monkeypatch.setattr( + native.tensor_parallel, "copy_to_tensor_model_parallel_region", lambda value: copied.append(value) or value + ) + monkeypatch.setattr(native.tensor_parallel, "scatter_to_sequence_parallel_region", lambda value: value.chunk(4)[1]) + monkeypatch.setattr(native.mpu, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(native.mpu, "get_tensor_model_parallel_rank", lambda: 1) + model = _injection_model(native, sequence_parallel=True) + image = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + _, mask, deepstack = native.Qwen3VLModel._inject_vision_embeddings( + model, torch.tensor([[10, 7, 10, 7, 7, 7, 7, 7]]), image, None, torch.tensor([[1, 2, 4]]), None + ) + assert len(copied) == 2 + assert mask.tolist() == [[True, False]] + assert torch.equal(deepstack[0], image[1:] * 2) + + +def test_missing_vision_is_not_silently_treated_as_text(native): + with pytest.raises(ValueError, match="pixel values"): + native.Qwen3VLModel._inject_vision_embeddings( + _injection_model(native), torch.tensor([[10]]), None, None, None, None + ) + + +def _vision_tp_worker(rank, rendezvous): + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + + from vime_plugins.models import qwen3_vl as native + + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=4, timeout=timedelta(seconds=45)) + try: + copy = native.tensor_parallel.copy_to_tensor_model_parallel_region + scatter = native.tensor_parallel.scatter_to_sequence_parallel_region + with ( + patch.object(torch.cuda, "current_device", lambda: torch.device("cpu")), + patch.object( + native.tensor_parallel, + "copy_to_tensor_model_parallel_region", + lambda x: copy(x, group=dist.group.WORLD), + ), + patch.object( + native.tensor_parallel, + "scatter_to_sequence_parallel_region", + lambda x: scatter(x, group=dist.group.WORLD), + ), + patch.object(native.mpu, "get_tensor_model_parallel_world_size", lambda: 4), + patch.object(native.mpu, "get_tensor_model_parallel_rank", lambda: rank), + ): + image = torch.arange(8, dtype=torch.float32).view(4, 2).requires_grad_() + embeddings, _, deepstack = native.Qwen3VLModel._inject_vision_embeddings( + _injection_model(native, sequence_parallel=True), + torch.tensor([[10, 7, 10, 7, 10, 7, 10, 7]]), + image, + None, + torch.tensor([[1, 4, 4]]), + None, + ) + (embeddings.sum() + sum(t.sum() for t in deepstack)).backward() + # All four replicas receive the same complete gradient, including + # image features consumed by other SP ranks (1 + 2 + 3 = 6). + torch.testing.assert_close(image.grad, torch.full_like(image, 6)) + finally: + dist.destroy_process_group() + + +@pytest.mark.integration +def test_trainable_vision_tp4_gradients_on_cpu(native, tmp_path): + torch.multiprocessing.spawn(_vision_tp_worker, args=((tmp_path / "gloo").as_uri(),), nprocs=4) + + +def test_deepstack_gradients_survive_main_recompute(native, monkeypatch): + from torch.utils.checkpoint import checkpoint + + from vime_plugins.models.qwen3_omni_transformer import Qwen3OmniTransformerBlock + + class Layer(torch.nn.Module): + def __init__(self, number): + super().__init__() + self.layer_number = number + + def forward(self, hidden_states, **kwargs): + return hidden_states * 2, None + + block = Qwen3OmniTransformerBlock.__new__(Qwen3OmniTransformerBlock) + torch.nn.Module.__init__(block) + block.config = SimpleNamespace( + fp8=False, distribute_saved_activations=False, recompute_method="uniform", recompute_num_layers=1 + ) + block.pre_process = True + block.layers = torch.nn.ModuleList([Layer(1), Layer(2)]) + block.num_layers_per_pipeline_rank = 2 + monkeypatch.setattr( + native.tensor_parallel, "checkpoint", lambda fn, distribute, *args: checkpoint(fn, *args, use_reentrant=True) + ) + hidden = torch.ones(4, 1, 2, requires_grad=True) + features = [torch.ones(1, 2, requires_grad=True), torch.ones(1, 2, requires_grad=True)] + output = block._checkpointed_forward( + hidden, + None, + None, + None, + None, + None, + None, + False, + visual_pos_masks=torch.tensor([[False, True, False, False]]), + deepstack_visual_embeds=features, + ) + output.sum().backward() + torch.testing.assert_close(hidden.grad, torch.full_like(hidden, 4)) + torch.testing.assert_close(features[0].grad, torch.full_like(features[0], 2)) + torch.testing.assert_close(features[1].grad, torch.ones_like(features[1])) + + +def test_interleaved_mrope_matches_hf_for_packed_images(native, hf_config): + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding + + from vime_plugins.models.qwen3_omni_moe import Qwen3OmniMultimodalRotaryEmbedding + + positions = native.build_packed_mrope_position_ids( + torch.tensor([[30, 10, 10, 10, 10, 7, 30, 10, 8]]), + [0, 6, 9], + torch.tensor([[1, 4, 4], [1, 2, 2]]), + None, + image_token_id=10, + video_token_id=20, + vision_start_token_id=30, + spatial_merge_size=2, + ) + assert positions[:, 0, 6].tolist() == [0, 0, 0] + hf_rope = Qwen3VLTextRotaryEmbedding(hf_config.text_config) + # Exercise the real main RoPE arithmetic without allocating a CUDA buffer. + rope = Qwen3OmniMultimodalRotaryEmbedding.__new__(Qwen3OmniMultimodalRotaryEmbedding) + torch.nn.Module.__init__(rope) + rope.inv_freq = hf_rope.inv_freq + rope.seq_len_interpolation_factor = None + rope.cp_group = None + rope.is_thd_format = True + freqs = rope(positions, [1, 1, 0])[:, 0, 0].unsqueeze(0) + cos, sin = hf_rope(torch.zeros(1, 9, 4), positions) + torch.testing.assert_close(freqs.cos(), cos) + torch.testing.assert_close(freqs.sin(), sin) + + +@pytest.mark.parametrize(("pp", "cp", "mtp"), [(2, 1, None), (1, 2, None), (1, 1, 1)]) +def test_provider_rejects_unimplemented_topologies(native, pp, cp, mtp): + with pytest.raises(ValueError, match="PP=1 and CP=1|MTP"): + native.get_qwen3_vl_model_provider( + SimpleNamespace(mtp_num_layers=mtp), + SimpleNamespace(pipeline_model_parallel_size=pp, context_parallel_size=cp), + None, + ) + + +def test_provider_uses_main_dense_deepstack_gpt(native, hf_config, monkeypatch): + calls = {} + + def gpt(**kwargs): + calls.update(kwargs) + return SimpleNamespace(share_embeddings_and_output_weights=False) + + monkeypatch.setattr(native, "Qwen3OmniMoeGPTModel", gpt) + monkeypatch.setattr(native, "_load_vision_model", lambda *args: torch.nn.Linear(8, 8)) + monkeypatch.setattr(native.AutoConfig, "from_pretrained", lambda *args, **kwargs: hf_config) + monkeypatch.setattr( + native, "get_gpt_layer_with_transformer_engine_spec", lambda *, qk_layernorm: {"qk_layernorm": qk_layernorm} + ) + args = SimpleNamespace( + hf_checkpoint="unused", + mtp_num_layers=None, + transformer_impl="transformer_engine", + normalization="RMSNorm", + padded_vocab_size=32, + max_position_embeddings=128, + fp16_lm_cross_entropy=False, + untie_embeddings_and_output_weights=True, + rotary_percent=1.0, + rotary_base=1000000, + ) + config = SimpleNamespace(pipeline_model_parallel_size=1, context_parallel_size=1, normalization="RMSNorm") + model = native.get_qwen3_vl_model_provider(args, config, None)() + assert calls["transformer_layer_spec"] == {"qk_layernorm": True} + assert calls["config"] is config + assert calls["config"].normalization == "RMSNorm" + assert calls["position_embedding_type"] == "mrope" + assert calls["rotary_base"] == 5000000 + assert calls["scatter_embedding_sequence_parallel"] is False + assert config.mrope_section == [1, 1, 0] + assert all(parameter.requires_grad for parameter in model.model.visual.parameters()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_read_file_slicing.py b/tests/test_read_file_slicing.py new file mode 100644 index 000000000..ad0cb362e --- /dev/null +++ b/tests/test_read_file_slicing.py @@ -0,0 +1,86 @@ +"""CPU unit tests for the ``path@[start:end]`` dataset-slicing syntax. + +``_parse_generalized_path``'s regex explicitly accepts a sign on both bounds +(``-?\\d*``), and the feature shipped with full negative-index support +(``df.iloc[row_slice]``). The streaming rewrite swapped that for +``itertools.islice``, which raises ``ValueError`` on any negative index — so +``@[-100:]`` ("the last 100 rows") went from working to crashing while the +parser still advertised it. + +Pinned here: non-negative slices keep streaming through ``islice``; slices +with a negative bound resolve against the real row count; parsing itself. +""" + +from __future__ import annotations + +import json + +import pytest + +from vime.utils.data import _parse_generalized_path, read_file + + +NUM_GPUS = 0 + +ROWS = [{"id": i} for i in range(10)] + + +@pytest.fixture +def jsonl_path(tmp_path): + path = tmp_path / "data.jsonl" + path.write_text("".join(json.dumps(row) + "\n" for row in ROWS)) + return str(path) + + +def _ids(generalized_path): + return [row["id"] for row in read_file(generalized_path)] + + +@pytest.mark.unit +def test_parse_generalized_path(): + assert _parse_generalized_path("/a/b.jsonl") == ("/a/b.jsonl", None) + assert _parse_generalized_path("/a/b.jsonl@[3:7]") == ("/a/b.jsonl", slice(3, 7)) + assert _parse_generalized_path("/a/b.jsonl@[-100:]") == ("/a/b.jsonl", slice(-100, None)) + assert _parse_generalized_path("/a/b.jsonl@[:-2]") == ("/a/b.jsonl", slice(None, -2)) + + +@pytest.mark.unit +def test_no_slice_reads_everything(jsonl_path): + assert _ids(jsonl_path) == list(range(10)) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "suffix,expected", + [ + ("@[0:3]", [0, 1, 2]), + ("@[3:]", [3, 4, 5, 6, 7, 8, 9]), + ("@[:4]", [0, 1, 2, 3]), + ], +) +def test_non_negative_slices(jsonl_path, suffix, expected): + assert _ids(jsonl_path + suffix) == expected + + +@pytest.mark.unit +@pytest.mark.parametrize( + "suffix,expected", + [ + ("@[-3:]", [7, 8, 9]), + ("@[:-2]", [0, 1, 2, 3, 4, 5, 6, 7]), + ("@[1:-1]", [1, 2, 3, 4, 5, 6, 7, 8]), + ("@[-5:-2]", [5, 6, 7]), + ], +) +def test_negative_slices(jsonl_path, suffix, expected): + assert _ids(jsonl_path + suffix) == expected + + +@pytest.mark.unit +def test_negative_slice_larger_than_file(jsonl_path): + # "@[-100:]" on a 10-row file is simply the whole file, like list slicing. + assert _ids(jsonl_path + "@[-100:]") == list(range(10)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_release_train.py b/tests/test_release_train.py new file mode 100644 index 000000000..a20caa0f9 --- /dev/null +++ b/tests/test_release_train.py @@ -0,0 +1,148 @@ +"""E2E smoke test for colocated ``--release-train``. + +The job runs two rollout steps so the actor group is released after each disk +weight update, then recreated from the saved Megatron checkpoint before the next +training step. +""" + +import os +import tempfile +from pathlib import Path +from shlex import quote + +import vime.utils.external_utils.command_utils as U + + +MODEL_NAME = "Qwen3.5-0.8B" +MODEL_TYPE = "qwen3.5-0.8B" +NUM_GPUS = 4 +NUM_ROLLOUT = 2 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) + + +def execute(): + with tempfile.TemporaryDirectory(prefix="vime_release_train_") as work_dir: + save_dir = Path(work_dir) / "mcore" + update_weight_dir = Path(work_dir) / "update_weight" + + ckpt_args = ( + f"--hf-checkpoint /root/models/{MODEL_NAME}/ " + f"--ref-load {TORCH_DIST_CKPT} " + "--release-train " + f"--save {quote(str(save_dir))} " + "--save-interval 1 " + ) + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + f"--num-rollout {NUM_ROLLOUT} " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 512 " + "--rollout-temperature 0.8 " + "--over-sampling-batch-size 8 " + "--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " + "--global-batch-size 16 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.01 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 1 " + "--vllm-gpu-memory-utilization 0.7 " + "--vllm-max-cudagraph-capture-size 16 " + ) + + disk_update_args = ( + "--update-weight-mode full " + "--update-weight-transport disk " + f"--update-weight-disk-dir {quote(str(update_weight_dir))} " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--loss-mask-type qwen3_5 " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {NUM_GPUS} " + "--colocate " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{vllm_args} " + f"{disk_update_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + latest_checkpoint = save_dir / "latest_checkpointed_iteration.txt" + assert latest_checkpoint.exists(), f"No Megatron checkpoint was saved under {save_dir}" + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_reloadable_process_group_memory_check.py b/tests/test_reloadable_process_group_memory_check.py index 94d9ebb43..dfe706b12 100644 --- a/tests/test_reloadable_process_group_memory_check.py +++ b/tests/test_reloadable_process_group_memory_check.py @@ -1,9 +1,13 @@ from __future__ import annotations +from datetime import timedelta + import pytest from vime.utils import reloadable_process_group as rpg +NUM_GPUS = 0 + @pytest.mark.unit def test_selected_comm_ops_skip_memory_check(): @@ -65,3 +69,139 @@ def fake_available_memory(): pass assert calls == ["available_memory"] + + +@pytest.mark.unit +def test_register_default_process_group_captures_rendezvous_state(monkeypatch): + timeout = timedelta(minutes=7) + monkeypatch.setattr(rpg, "default_process_group_states", {}) + monkeypatch.setattr(rpg.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rpg.dist, "get_backend", lambda: "nccl") + monkeypatch.setattr(rpg.dist, "get_rank", lambda: 3) + monkeypatch.setattr(rpg.dist, "get_world_size", lambda: 8) + monkeypatch.setattr(rpg, "_get_default_store", lambda: "rendezvous-store") + + rpg.register_default_process_group(timeout=timeout) + + state = rpg.default_process_group_states[rpg.os.getpid()] + assert state.backend == "nccl" + assert state.timeout == timeout + assert state.store == "rendezvous-store" + assert state.rank == 3 + assert state.world_size == 8 + assert not state.accelerator_world_destroyed + + +@pytest.mark.unit +def test_world_and_subgroups_follow_destroy_reload_order(monkeypatch): + timeout = timedelta(minutes=2) + state = rpg._DefaultProcessGroupState( + backend="nccl", + timeout=timeout, + store="base-store", + rank=1, + world_size=4, + ) + monkeypatch.setattr(rpg, "default_process_group_states", {rpg.os.getpid(): state}) + + events = [] + + def barrier(group=None): + events.append(("barrier", "WORLD" if group is None else group)) + + def init_process_group(**kwargs): + events.append(("init", kwargs)) + + monkeypatch.setattr(rpg.dist, "barrier", barrier) + monkeypatch.setattr(rpg.dist, "destroy_process_group", lambda: events.append(("destroy_world",))) + monkeypatch.setattr(rpg.dist, "init_process_group", init_process_group) + monkeypatch.setattr(rpg, "PrefixStore", lambda prefix, store: (prefix, store)) + monkeypatch.setattr(rpg, "get_gloo_group", lambda: "canonical-gloo") + monkeypatch.setattr(rpg, "set_gloo_group", lambda group: events.append(("set_gloo", group))) + monkeypatch.setattr(rpg, "_get_default_group", lambda: "cpu-world") + monkeypatch.setattr(rpg, "init_gloo_group", lambda: events.append(("init_canonical_gloo",))) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "invalidate_process_groups", + staticmethod(lambda: events.append(("invalidate_subgroups",))), + ) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "reload_process_groups", + staticmethod(lambda: events.append(("reload_subgroups",))), + ) + + rpg.destroy_process_groups() + + assert state.accelerator_world_destroyed + assert state.generation == 1 + assert events == [ + ("barrier", "canonical-gloo"), + ("destroy_world",), + ("invalidate_subgroups",), + ("set_gloo", None), + ( + "init", + { + "backend": "gloo", + "store": ("vime-reloadable-world-1-gloo", "base-store"), + "rank": 1, + "world_size": 4, + "timeout": timeout, + }, + ), + ("set_gloo", "cpu-world"), + ] + + events.clear() + rpg.reload_process_groups() + + assert not state.accelerator_world_destroyed + assert state.generation == 2 + assert events == [ + ("barrier", "WORLD"), + ("destroy_world",), + ("set_gloo", None), + ( + "init", + { + "backend": "nccl", + "store": ("vime-reloadable-world-2-nccl", "base-store"), + "rank": 1, + "world_size": 4, + "timeout": timeout, + }, + ), + ("init_canonical_gloo",), + ("reload_subgroups",), + ] + + +@pytest.mark.unit +def test_unregistered_world_preserves_subgroup_only_behavior(monkeypatch): + events = [] + monkeypatch.setattr(rpg, "default_process_group_states", {}) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "destroy_process_groups", + staticmethod(lambda: events.append("destroy_subgroups")), + ) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "reload_process_groups", + staticmethod(lambda: events.append("reload_subgroups")), + ) + monkeypatch.setattr( + rpg.dist, + "destroy_process_group", + lambda: pytest.fail("unregistered WORLD must not be destroyed"), + ) + + rpg.destroy_process_groups() + rpg.reload_process_groups() + + assert events == ["destroy_subgroups", "reload_subgroups"] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_reloadable_process_group_world.py b/tests/test_reloadable_process_group_world.py new file mode 100644 index 000000000..dce01abab --- /dev/null +++ b/tests/test_reloadable_process_group_world.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch.distributed as dist +import torch.multiprocessing as mp + +from vime.utils import distributed_utils +from vime.utils import reloadable_process_group as rpg + +NUM_GPUS = 0 + + +def _run_pp_group_reload_worker(rank: int, world_size: int, rendezvous_path: str) -> None: + timeout = timedelta(seconds=30) + rpg.monkey_patch_torch_dist() + dist.init_process_group( + backend="gloo", + init_method=f"file://{rendezvous_path}", + rank=rank, + world_size=world_size, + timeout=timeout, + ) + distributed_utils.init_gloo_group() + rpg.register_default_process_group(timeout=timeout) + + # Exercise the accelerator lifecycle with Gloo so this remains a CPU test. + # The contract under test is WORLD/subgroup teardown ordering, not a vendor + # communication backend. + rpg._uses_accelerator_backend = lambda _backend: True + + group_specs = [ + ([0], "TP_0"), + ([1], "TP_1"), + ([2], "TP_2"), + ([3], "TP_3"), + ([0, 1, 2, 3], "PP"), + ([0, 3], "EMBEDDING"), + ([0], "POSITION_EMBEDDING"), + ([0, 2], "DP_0"), + ([1, 3], "DP_1"), + ] + groups = [ + dist.new_group(ranks=ranks, backend="gloo", timeout=timeout, group_desc=desc) for ranks, desc in group_specs + ] + pp_group = groups[4] + + for generation in range(2): + rpg.destroy_process_groups() + assert all(group.group is None for group in groups) + + rpg.reload_process_groups() + assert all(group.group is not None for group in groups) + + # Keep this NUM_GPUS=0 regression independent of CUDA-specific memory + # checks while still exercising a real collective on the reloaded PP group. + dist.barrier(group=pp_group) + assert dist.get_world_size(pp_group) == world_size + + state = rpg.default_process_group_states[rpg.os.getpid()] + assert state.generation == 2 * (generation + 1) + + dist.destroy_process_group() + + +def _run_backend_normalization_worker(_rank: int) -> None: + calls = [] + mapped_backends = [] + + def old_new_group(*args, **kwargs): + calls.append((args, kwargs)) + return f"group-{len(calls)}" + + def process_group_backend(backend): + mapped_backends.append(backend) + return "mccl" if backend == "nccl" else backend + + rpg.old_new_group_dict.clear() + rpg.default_process_group_states.clear() + rpg.dist.new_group = old_new_group + rpg.dist.get_backend = lambda: "gloo" + rpg.accelerator.process_group_backend = process_group_backend + rpg.monkey_patch_torch_dist() + + gloo_group = rpg.dist.new_group(ranks=[0], backend="gloo") + assert gloo_group == "group-1" + assert calls[-1][1]["backend"] == "gloo" + assert mapped_backends == [] + + mccl_group = rpg.dist.new_group([0], None, "nccl") + assert mccl_group == "group-2" + assert calls[-1][0][2] == "mccl" + assert mapped_backends == ["nccl"] + + +@pytest.mark.unit +@pytest.mark.parametrize("backend", ["nccl", "mccl", "cpu:gloo,musa:mccl"]) +def test_accelerator_backend_detection(backend): + assert rpg._uses_accelerator_backend(backend) + + +@pytest.mark.unit +def test_new_group_normalizes_only_logical_nccl_backend(): + # monkey_patch_torch_dist replaces process-wide torch.distributed symbols, + # so isolate this behavior check in a spawned process. + mp.spawn(_run_backend_normalization_worker, nprocs=1, join=True) + + +@pytest.mark.unit +def test_register_default_process_group_captures_rendezvous_state(monkeypatch): + timeout = timedelta(minutes=7) + monkeypatch.setattr(rpg, "default_process_group_states", {}) + monkeypatch.setattr(rpg.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rpg.dist, "get_backend", lambda: "nccl") + monkeypatch.setattr(rpg.dist, "get_rank", lambda: 3) + monkeypatch.setattr(rpg.dist, "get_world_size", lambda: 8) + monkeypatch.setattr(rpg, "_get_default_store", lambda: "rendezvous-store") + + rpg.register_default_process_group(timeout=timeout) + + state = rpg.default_process_group_states[rpg.os.getpid()] + assert state.backend == "nccl" + assert state.timeout == timeout + assert state.store == "rendezvous-store" + assert state.rank == 3 + assert state.world_size == 8 + assert not state.accelerator_world_destroyed + + +@pytest.mark.unit +def test_world_and_subgroups_follow_destroy_reload_order(monkeypatch): + timeout = timedelta(minutes=2) + state = rpg._DefaultProcessGroupState( + backend="nccl", + timeout=timeout, + store="base-store", + rank=1, + world_size=4, + ) + monkeypatch.setattr(rpg, "default_process_group_states", {rpg.os.getpid(): state}) + + events = [] + + def barrier(group=None): + events.append(("barrier", "WORLD" if group is None else group)) + + def init_process_group(**kwargs): + events.append(("init", kwargs)) + + monkeypatch.setattr(rpg.dist, "barrier", barrier) + monkeypatch.setattr(rpg.dist, "destroy_process_group", lambda: events.append(("destroy_world",))) + monkeypatch.setattr(rpg.dist, "init_process_group", init_process_group) + monkeypatch.setattr(rpg, "PrefixStore", lambda prefix, store: (prefix, store)) + monkeypatch.setattr(rpg, "get_gloo_group", lambda: "canonical-gloo") + monkeypatch.setattr(rpg, "set_gloo_group", lambda group: events.append(("set_gloo", group))) + monkeypatch.setattr(rpg, "_get_default_group", lambda: "cpu-world") + monkeypatch.setattr(rpg, "init_gloo_group", lambda: events.append(("init_canonical_gloo",))) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "invalidate_process_groups", + staticmethod(lambda: events.append(("invalidate_subgroups",))), + ) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "reload_process_groups", + staticmethod(lambda: events.append(("reload_subgroups",))), + ) + + rpg.destroy_process_groups() + + assert state.accelerator_world_destroyed + assert state.generation == 1 + assert events == [ + ("barrier", "canonical-gloo"), + ("destroy_world",), + ("invalidate_subgroups",), + ("set_gloo", None), + ( + "init", + { + "backend": "gloo", + "store": ("vime-reloadable-world-1-gloo", "base-store"), + "rank": 1, + "world_size": 4, + "timeout": timeout, + }, + ), + ("set_gloo", "cpu-world"), + ] + + events.clear() + rpg.reload_process_groups() + + assert not state.accelerator_world_destroyed + assert state.generation == 2 + assert events == [ + ("barrier", "WORLD"), + ("destroy_world",), + ("set_gloo", None), + ( + "init", + { + "backend": "nccl", + "store": ("vime-reloadable-world-2-nccl", "base-store"), + "rank": 1, + "world_size": 4, + "timeout": timeout, + }, + ), + ("init_canonical_gloo",), + ("reload_subgroups",), + ] + + +@pytest.mark.unit +def test_invalidating_wrappers_drops_every_stale_group_handle(monkeypatch): + groups = [SimpleNamespace(group=object()), SimpleNamespace(group=object())] + monkeypatch.setattr(rpg.ReloadableProcessGroup, "GROUPS", {rpg.os.getpid(): groups}) + + rpg.ReloadableProcessGroup.invalidate_process_groups() + + assert all(group.group is None for group in groups) + + +@pytest.mark.unit +def test_pp_topology_survives_repeated_world_and_subgroup_reload(tmp_path): + world_size = 4 + mp.spawn( + _run_pp_group_reload_worker, + args=(world_size, str(tmp_path / "rendezvous")), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.unit +def test_unregistered_world_preserves_subgroup_only_behavior(monkeypatch): + events = [] + monkeypatch.setattr(rpg, "default_process_group_states", {}) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "destroy_process_groups", + staticmethod(lambda: events.append("destroy_subgroups")), + ) + monkeypatch.setattr( + rpg.ReloadableProcessGroup, + "reload_process_groups", + staticmethod(lambda: events.append("reload_subgroups")), + ) + monkeypatch.setattr( + rpg.dist, + "destroy_process_group", + lambda: pytest.fail("unregistered WORLD must not be destroyed"), + ) + + rpg.destroy_process_groups() + rpg.reload_process_groups() + + assert events == ["destroy_subgroups", "reload_subgroups"] + + +@pytest.mark.unit +def test_torch_213_collectives_forward_to_inner_group(): + calls = [] + + class Recorder: + def _fwd(self, method, *args, **kwargs): + calls.append((method, args, kwargs)) + + recorder = Recorder() + methods = [ + "all_gather_single", + "all_gather_single_coalesced", + "reduce_scatter_single", + "reduce_scatter_single_coalesced", + "monitored_barrier", + ] + for method in methods: + getattr(rpg.ReloadableProcessGroup, method)(recorder, "tensor", async_op=True) + + assert calls == [(method, ("tensor",), {"async_op": True}) for method in methods] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rollout_data_utils.py b/tests/test_rollout_data_utils.py new file mode 100644 index 000000000..9a22fd5cb --- /dev/null +++ b/tests/test_rollout_data_utils.py @@ -0,0 +1,102 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vime.observability.rollout_data_utils import ( + load_debug_rollout_data, + save_debug_rollout_data, + tensorize_rollout_data_for_training, + validate_rollout_routed_experts_for_replay, +) +from vime.utils.types import Sample + +NUM_GPUS = 0 + + +def _args(): + return SimpleNamespace( + num_layers=6, + moe_router_topk=2, + moe_layer_freq=[0, 0, 0, 1, 1, 1], + ) + + +def test_r3_validation_accepts_dense_zeros_and_complete_moe_routes(): + routes = torch.zeros((4, 6, 2), dtype=torch.uint8) + routes[:, 3:, 1] = 7 + validate_rollout_routed_experts_for_replay([routes], _args()) + + +def test_r3_validation_rejects_missing_pipeline_layers(): + routes = torch.zeros((4, 6, 2), dtype=torch.uint8) + routes[:, 3, 1] = 7 + + with pytest.raises(ValueError, match=r"all zero.*\[4, 5\]"): + validate_rollout_routed_experts_for_replay([routes], _args()) + + +def test_r3_validation_rejects_wrong_shape(): + routes = torch.zeros((4, 5, 2), dtype=torch.uint8) + + with pytest.raises(ValueError, match="Invalid rollout routed-experts shape"): + validate_rollout_routed_experts_for_replay([routes], _args()) + + +def test_tensorize_rollout_data_for_training_normalizes_cpu_tensors(): + readonly_tokens = np.array([1, 2, 3]) + readonly_tokens.flags.writeable = False + rollout_data = { + "tokens": [readonly_tokens], + "loss_masks": [[1, 0]], + "multimodal_train_inputs": [ + { + "pixel_values": torch.tensor([1.0], requires_grad=True), + "metadata": "unchanged", + } + ], + "rollout_mask_sums": [2], + } + + tensorize_rollout_data_for_training(rollout_data) + + assert rollout_data["tokens"][0].dtype == torch.long + assert rollout_data["loss_masks"][0].dtype == torch.int + assert rollout_data["multimodal_train_inputs"][0]["metadata"] == "unchanged" + assert not rollout_data["multimodal_train_inputs"][0]["pixel_values"].requires_grad + assert rollout_data["rollout_mask_sums"].dtype == torch.float32 + + +def test_save_and_load_debug_rollout_data_round_trip(tmp_path): + path_template = str(tmp_path / "rollout_{rollout_id}.pt") + samples = [ + Sample(index=1, rollout_id=3, prompt="question", response="answer", response_length=1), + ] + + save_debug_rollout_data(path_template, samples, rollout_id=3, evaluation=False) + loaded = load_debug_rollout_data(path_template, rollout_id=3) + + assert len(loaded) == 1 + assert loaded[0].index == 1 + assert loaded[0].rollout_id == 3 + assert loaded[0].prompt == "question" + assert loaded[0].response == "answer" + + +def test_save_debug_eval_rollout_data_flattens_datasets(tmp_path): + path_template = str(tmp_path / "rollout_{rollout_id}.pt") + data = { + "math": {"samples": [Sample(index=1)]}, + "code": {"samples": [Sample(index=2)]}, + } + + save_debug_rollout_data(path_template, data, rollout_id=4, evaluation=True) + + saved = torch.load(tmp_path / "rollout_eval_4.pt", weights_only=False) + assert saved["rollout_id"] == 4 + assert [sample["index"] for sample in saved["samples"]] == [1, 2] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rollout_metrics.py b/tests/test_rollout_metrics.py new file mode 100644 index 000000000..273802d67 --- /dev/null +++ b/tests/test_rollout_metrics.py @@ -0,0 +1,221 @@ +import base64 +from argparse import Namespace + +import numpy as np +import pytest +import torch + +from vime.observability.rollout_metrics import _compute_top_p_kept_vocab_metrics +from vime.utils.misc import decode_int32_meta_array +from vime.utils.types import Sample + +NUM_GPUS = 0 + + +def _make_args(): + return Namespace(vllm_speculative_algorithm=False, num_layers=2, moe_router_topk=2) + + +@pytest.mark.unit +def test_top_p_kept_vocab_metric_uses_loss_mask(): + samples = [ + Sample( + response_length=4, + loss_mask=torch.tensor([1, 0, 1, 0], dtype=torch.int32), + rollout_top_p_token_offsets=torch.tensor([0, 3, 8, 10, 20], dtype=torch.int32), + ), + Sample( + response_length=2, + loss_mask=None, + rollout_top_p_token_offsets=torch.tensor([0, 4, 9], dtype=torch.int32), + ), + ] + + metrics = _compute_top_p_kept_vocab_metrics(samples) + + assert metrics["top_p_kept_vocab_per_token"] == pytest.approx(3.5) + + +@pytest.mark.unit +def test_top_p_kept_vocab_metric_skips_removed_samples(): + samples = [ + Sample( + response_length=3, + loss_mask=[1, 1, 1], + remove_sample=True, + rollout_top_p_token_offsets=torch.tensor([0, 2, 4, 6], dtype=torch.int32), + ) + ] + + assert _compute_top_p_kept_vocab_metrics(samples) == {} + + +def _b64_int32(values: list[int]) -> str: + return base64.b64encode(np.array(values, dtype=np.int32).tobytes()).decode("ascii") + + +@pytest.mark.unit +def test_decode_int32_meta_array_decodes_base64_to_tensor(): + decoded = decode_int32_meta_array({"routed_experts": _b64_int32([1, 2, 3])}, "routed_experts") + + assert torch.is_tensor(decoded) + assert decoded.dtype == torch.int32 + torch.testing.assert_close(decoded, torch.tensor([1, 2, 3], dtype=torch.int32)) + + +@pytest.mark.unit +def test_append_response_tokens_merges_top_p_tensors(): + sample = Sample( + tokens=[0, 1], + response_length=1, + loss_mask=[1], + rollout_log_probs=[-0.3], + rollout_top_p_token_ids=torch.tensor([1], dtype=torch.int32), + rollout_top_p_token_offsets=torch.tensor([0, 1], dtype=torch.int32), + ) + + sample.append_response_tokens( + _make_args(), + tokens=[10, 20], + log_probs=[-0.1, -0.2], + trainable=True, + meta_info={ + "top_p_token_ids": _b64_int32([10, 11, 20]), + "top_p_token_offsets": _b64_int32([0, 2, 3]), + "finish_reason": {"type": "stop"}, + }, + ) + + assert sample.tokens == [0, 1, 10, 20] + assert sample.response_length == 3 + assert sample.loss_mask == [1, 1, 1] + assert sample.rollout_log_probs == [-0.3, -0.1, -0.2] + torch.testing.assert_close(sample.rollout_top_p_token_ids, torch.tensor([1, 10, 11, 20], dtype=torch.int32)) + torch.testing.assert_close(sample.rollout_top_p_token_offsets, torch.tensor([0, 1, 3, 4], dtype=torch.int32)) + + +@pytest.mark.unit +def test_append_response_tokens_can_skip_terminal_status_for_streaming_chunks(): + sample = Sample( + tokens=[0, 1], + response_length=1, + loss_mask=[1], + rollout_log_probs=[-0.3], + rollout_top_p_token_ids=torch.tensor([1], dtype=torch.int32), + rollout_top_p_token_offsets=torch.tensor([0, 1], dtype=torch.int32), + ) + + sample.append_response_tokens( + _make_args(), + tokens=[10, 20], + log_probs=[-0.1, -0.2], + trainable=True, + meta_info={ + "top_p_token_ids": _b64_int32([10, 11, 20]), + "top_p_token_offsets": _b64_int32([0, 2, 3]), + "finish_reason": {"type": "stop"}, + }, + update_terminal_info=False, + ) + + assert sample.status is Sample.Status.PENDING + assert sample.loss_mask == [1, 1, 1] + assert sample.rollout_log_probs == [-0.3, -0.1, -0.2] + torch.testing.assert_close(sample.rollout_top_p_token_ids, torch.tensor([1, 10, 11, 20], dtype=torch.int32)) + torch.testing.assert_close(sample.rollout_top_p_token_offsets, torch.tensor([0, 1, 3, 4], dtype=torch.int32)) + + +@pytest.mark.unit +def test_append_response_tokens_decodes_routed_experts(): + sample = Sample(tokens=[101, 102, 103]) + + sample.append_response_tokens( + _make_args(), + tokens=[], + trainable=True, + meta_info={ + "routed_experts": _b64_int32([0, 1, 2, 3, 4, 5, 6, 7]), + "finish_reason": {"type": "stop"}, + }, + ) + + assert sample.rollout_routed_experts.shape == (2, 2, 2) + torch.testing.assert_close( + sample.rollout_routed_experts, + torch.tensor([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=torch.int32), + ) + + +@pytest.mark.unit +def test_append_response_tokens_ignores_split_pd_routed_experts(): + sample = Sample(tokens=[101, 102, 103, 104]) + + sample.append_response_tokens( + _make_args(), + tokens=[], + trainable=True, + meta_info={ + "pd_prefill_routed_experts": _b64_int32([0, 1, 2, 3, 4, 5, 6, 7]), + "pd_decode_routed_experts": _b64_int32([8, 9, 10, 11]), + "finish_reason": {"type": "stop"}, + }, + ) + + assert sample.rollout_routed_experts is None + + +@pytest.mark.unit +def test_append_response_tokens_rejects_mismatched_routed_experts_shape(): + sample = Sample(tokens=[101, 102, 103]) + + with pytest.raises(ValueError, match="routed_experts element count"): + sample.append_response_tokens( + _make_args(), + tokens=[], + trainable=True, + meta_info={ + "routed_experts": _b64_int32([0, 1, 2, 3]), + "finish_reason": {"type": "stop"}, + }, + ) + + +@pytest.mark.unit +def test_append_response_tokens_pads_top_p_for_non_trainable_tokens(): + sample = Sample( + tokens=[0, 1], + response_length=1, + loss_mask=[1], + rollout_log_probs=[-0.1], + rollout_top_p_token_ids=torch.tensor([10, 11], dtype=torch.int32), + rollout_top_p_token_offsets=torch.tensor([0, 2], dtype=torch.int32), + ) + + sample.append_response_tokens(tokens=[200, 201, 202], trainable=False) + + assert sample.tokens == [0, 1, 200, 201, 202] + assert sample.response_length == 4 + assert sample.loss_mask == [1, 0, 0, 0] + assert sample.rollout_log_probs == [-0.1, 0.0, 0.0, 0.0] + torch.testing.assert_close(sample.rollout_top_p_token_ids, torch.tensor([10, 11], dtype=torch.int32)) + torch.testing.assert_close(sample.rollout_top_p_token_offsets, torch.tensor([0, 2, 2, 2, 2], dtype=torch.int32)) + + +@pytest.mark.unit +def test_append_response_tokens_requires_trainable_log_probs(): + sample = Sample() + + with pytest.raises(ValueError, match="trainable response tokens require rollout log probabilities"): + sample.append_response_tokens(tokens=[10], trainable=True) + + +@pytest.mark.unit +def test_append_response_tokens_rejects_non_trainable_log_probs(): + sample = Sample() + + with pytest.raises(ValueError, match="non-trainable response tokens should not pass rollout log probabilities"): + sample.append_response_tokens(tokens=[10], log_probs=[-0.1], trainable=False) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rollout_sample_hooks.py b/tests/test_rollout_sample_hooks.py new file mode 100644 index 000000000..6ee145227 --- /dev/null +++ b/tests/test_rollout_sample_hooks.py @@ -0,0 +1,59 @@ +import asyncio +import types + +import pytest + +from vime.rollout.sample_hooks import apply_rollout_sample_hooks, set_current_rollout_id +from vime.utils.types import Sample + +NUM_GPUS = 0 + + +def sync_hook(args, sample, *, rollout_id=None): + sample.metadata["sync_hook"] = (args.marker, rollout_id) + + +async def async_hook(args, sample, *, evaluation=False): + sample.metadata["async_hook"] = evaluation + return sample + + +def invalid_hook(args, sample): + return {"sample": sample} + + +@pytest.mark.unit +def test_rollout_sample_hooks_preserve_nested_shape_and_filter_kwargs(): + args = types.SimpleNamespace( + marker="seen", + rollout_sample_hook_path=[f"{__name__}.sync_hook", f"{__name__}.async_hook"], + ) + samples = [[Sample(index=0, metadata={})], [Sample(index=1, metadata={})]] + set_current_rollout_id(7) + + result = asyncio.run(apply_rollout_sample_hooks(args, samples, evaluation=True, ignored="value")) + + assert result is not samples + assert [[sample.index for sample in group] for group in result] == [[0], [1]] + for group in result: + assert group[0].metadata == {"sync_hook": ("seen", 7), "async_hook": True} + + +@pytest.mark.unit +def test_rollout_sample_hook_rejects_invalid_return_type(): + args = types.SimpleNamespace(rollout_sample_hook_path=[f"{__name__}.invalid_hook"]) + + with pytest.raises(TypeError, match="expected Sample or None"): + asyncio.run(apply_rollout_sample_hooks(args, Sample(index=0))) + + +@pytest.mark.unit +def test_rollout_sample_hooks_are_noop_when_unconfigured(): + args = types.SimpleNamespace(rollout_sample_hook_path=[]) + sample = Sample(index=0) + + assert asyncio.run(apply_rollout_sample_hooks(args, sample)) is sample + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rollout_validation.py b/tests/test_rollout_validation.py deleted file mode 100644 index 4e3d63794..000000000 --- a/tests/test_rollout_validation.py +++ /dev/null @@ -1,63 +0,0 @@ -import pytest - -from vime.ray.rollout_validation import validate_server_group_gpu_indices - - -NUM_GPUS = 0 - - -@pytest.mark.unit -def test_validate_server_group_gpu_indices_accepts_valid_config(): - validate_server_group_gpu_indices( - worker_type="regular", - gpu_offset=2, - num_gpus_per_engine=1, - num_gpu_per_engine=1, - num_engines=2, - num_available_gpus=4, - rollout_num_gpus=4, - rollout_num_gpus_per_engine=1, - ) - - -@pytest.mark.unit -def test_validate_server_group_gpu_indices_allows_empty_group(): - validate_server_group_gpu_indices( - worker_type="placeholder", - gpu_offset=4, - num_gpus_per_engine=1, - num_gpu_per_engine=1, - num_engines=0, - num_available_gpus=4, - rollout_num_gpus=4, - rollout_num_gpus_per_engine=1, - ) - - -@pytest.mark.unit -def test_validate_server_group_gpu_indices_reports_config_context(): - with pytest.raises(ValueError) as exc_info: - validate_server_group_gpu_indices( - worker_type="regular", - gpu_offset=3, - num_gpus_per_engine=2, - num_gpu_per_engine=2, - num_engines=1, - num_available_gpus=4, - rollout_num_gpus=4, - rollout_num_gpus_per_engine=2, - ) - - message = str(exc_info.value) - assert "worker_type=regular" in message - assert "gpu_offset=3" in message - assert "num_gpus_per_engine=2" in message - assert "num_engines=1" in message - assert "required_gpu_slots=5" in message - assert "len(reordered_gpu_ids)=4" in message - assert "rollout_num_gpus=4" in message - assert "rollout_num_gpus_per_engine=2" in message - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_sample.py b/tests/test_sample.py index bf9e22d77..bc9d3ff6d 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -8,7 +8,7 @@ sample loses its status / spec_info / prefix_cache_info on the way to the trainer with no crash signal. - 2. ``update_from_meta_info`` finish_reason → Status enum mapping + 2. ``append_response_tokens`` finish_reason → Status enum mapping (length→TRUNCATED, abort→ABORTED, stop→COMPLETED). The match statement at types.py:176-182 is the only place vLLM's finish_reason gets translated; a typo'd enum or removed case here @@ -50,6 +50,8 @@ def _make_sample(**overrides) -> Sample: loss_mask=[1, 1, 0, 1, 1], weight_versions=["v1"], rollout_log_probs=[-0.1, -0.2], + rollout_top_p_token_ids=[10, 11, 12, 20], + rollout_top_p_token_offsets=[0, 3, 4], rollout_routed_experts=[[0, 1], [2, 3]], remove_sample=False, teacher_log_probs=[-0.3, -0.4], @@ -131,6 +133,8 @@ def test_round_trip_preserves_every_field(): "loss_mask", "weight_versions", "rollout_log_probs", + "rollout_top_p_token_ids", + "rollout_top_p_token_offsets", "rollout_routed_experts", "remove_sample", "teacher_log_probs", @@ -169,12 +173,12 @@ def test_round_trip_through_default_constructed_sample(): # --------------------------------------------------------------------------- -# update_from_meta_info — finish_reason → Status mapping +# append_response_tokens — finish_reason → Status mapping # --------------------------------------------------------------------------- def _make_args(speculative: bool = False) -> argparse.Namespace: - """``update_from_meta_info`` only consults ``args.vllm_speculative_config`` + """``append_response_tokens`` only consults ``args.vllm_speculative_algorithm`` — minimal stub is enough.""" return argparse.Namespace(vllm_speculative_config=speculative) @@ -193,8 +197,10 @@ def test_status_mapping_for_each_finish_reason(finish_reason, expected_status): finish_reason ever gets translated. Each branch must hit the right enum; a typo in the enum name would crash later in unrelated places.""" sample = Sample() - sample.update_from_meta_info( + sample.append_response_tokens( _make_args(), + tokens=[], + trainable=True, meta_info={"finish_reason": {"type": finish_reason}}, ) assert sample.status is expected_status @@ -207,8 +213,10 @@ def test_unknown_finish_reason_leaves_status_unchanged(): a default doesn't silently break this contract.""" sample = Sample() sample.status = Sample.Status.PENDING - sample.update_from_meta_info( + sample.append_response_tokens( _make_args(), + tokens=[], + trainable=True, meta_info={"finish_reason": {"type": "something_new"}}, ) assert sample.status is Sample.Status.PENDING @@ -221,8 +229,10 @@ def test_weight_version_is_appended_when_present(): version produced each chunk.""" sample = Sample() sample.weight_versions = ["v1"] - sample.update_from_meta_info( + sample.append_response_tokens( _make_args(), + tokens=[], + trainable=True, meta_info={ "finish_reason": {"type": "stop"}, "weight_version": "v2", @@ -233,13 +243,15 @@ def test_weight_version_is_appended_when_present(): @pytest.mark.unit def test_prefix_cache_info_is_accumulated_across_calls(): - """Every call to update_from_meta_info adds to prefix_cache_info + """Every call to append_response_tokens with terminal metadata adds to prefix_cache_info (types.py:171). Multi-turn rollouts call this once per turn — the counts must accumulate, not overwrite.""" sample = Sample() for prompt_tokens, cached_tokens in [(100, 0), (200, 50)]: - sample.update_from_meta_info( + sample.append_response_tokens( _make_args(), + tokens=[], + trainable=True, meta_info={ "finish_reason": {"type": "stop"}, "prompt_tokens": prompt_tokens, @@ -262,11 +274,11 @@ def test_spec_info_only_updated_when_speculative_enabled(): } no_spec = Sample() - no_spec.update_from_meta_info(_make_args(speculative=False), meta_info=meta_info) + no_spec.append_response_tokens(_make_args(speculative=False), tokens=[], trainable=True, meta_info=meta_info) assert no_spec.spec_info.spec_accept_token_num == 0 with_spec = Sample() - with_spec.update_from_meta_info(_make_args(speculative=True), meta_info=meta_info) + with_spec.append_response_tokens(_make_args(speculative=True), tokens=[], trainable=True, meta_info=meta_info) assert with_spec.spec_info.spec_accept_token_num == 7 assert with_spec.spec_info.spec_draft_token_num == 10 diff --git a/tests/test_stateless_adam.py b/tests/test_stateless_adam.py new file mode 100644 index 000000000..359e3c312 --- /dev/null +++ b/tests/test_stateless_adam.py @@ -0,0 +1,62 @@ +import pytest +import torch + +from vime.backends.megatron_utils.stateless_adam import StatelessAdam + +NUM_GPUS = 0 + + +def _run_with_reinitialized_adam(param, grads, *, adam_w_mode): + param = param.clone().detach() + optimizer_cls = torch.optim.AdamW if adam_w_mode else torch.optim.Adam + for grad in grads: + param = param.detach().requires_grad_(True) + optimizer = optimizer_cls( + [param], + lr=0.03, + betas=(0.9, 0.98), + eps=1e-6, + weight_decay=0.1, + ) + param.grad = grad.clone() + optimizer.step() + param = param.detach() + return param + + +@pytest.mark.unit +@pytest.mark.parametrize("adam_w_mode", [True, False]) +def test_stateless_adam_matches_reinitialized_adam_each_step(adam_w_mode): + torch.manual_seed(0) + initial_param = torch.randn(8, dtype=torch.float64) + grads = [torch.randn_like(initial_param) for _ in range(4)] + param = initial_param.clone() + optimizer = StatelessAdam( + [param], + lr=0.03, + betas=(0.9, 0.98), + eps=1e-6, + weight_decay=0.1, + adam_w_mode=adam_w_mode, + ) + + for grad in grads: + param.grad = grad.clone() + optimizer.step() + optimizer.zero_grad() + + expected = _run_with_reinitialized_adam(initial_param, grads, adam_w_mode=adam_w_mode) + torch.testing.assert_close(param, expected) + + +@pytest.mark.unit +def test_stateless_adam_does_not_persist_moment_tensors(): + param = torch.tensor([1.0, -2.0]) + optimizer = StatelessAdam([param]) + + assert optimizer.state == {} + assert optimizer.state_dict()["state"] == {} + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py new file mode 100644 index 000000000..4e677ef27 --- /dev/null +++ b/tests/test_tau_bench_token_delta.py @@ -0,0 +1,213 @@ +import importlib.util +import re +from pathlib import Path + +import pytest + +NUM_GPUS = 0 + + +TOKEN_DELTA_PATH = Path(__file__).parents[1] / "examples" / "tau-bench" / "token_delta.py" + + +def _load_get_token_delta(): + if not TOKEN_DELTA_PATH.exists(): + pytest.fail("tau-bench token_delta helper does not exist") + + spec = importlib.util.spec_from_file_location("tau_bench_token_delta", TOKEN_DELTA_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module.get_token_delta + + +class HistoryRewritingTokenizer: + """Minimal chat template that hides old reasoning after a new user turn.""" + + @staticmethod + def _strip_reasoning(content: str) -> str: + return re.sub(r".*?", "", content, flags=re.DOTALL) + + def apply_chat_template(self, messages, *, add_generation_prompt, tokenize): + assert tokenize is False + last_user = max((i for i, message in enumerate(messages) if message["role"] == "user"), default=-1) + rendered = [] + for i, message in enumerate(messages): + content = message["content"] + if message["role"] == "assistant" and i < last_user: + content = self._strip_reasoning(content) + rendered.append(f'<{message["role"]}>{content}') + if add_generation_prompt: + rendered.append("") + return "".join(rendered) + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + return list(text.encode()) + + @staticmethod + def decode(token_ids): + return bytes(token_ids).decode() + + +class BoundaryMergingTokenizer(HistoryRewritingTokenizer): + """Tokenizer where the generation-prefix tail merges with a leading newline.""" + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + raw = text.encode() + token_ids = [] + index = 0 + while index < len(raw): + if raw[index : index + 2] == b">\n": + token_ids.append(1000) + index += 2 + else: + token_ids.append(raw[index]) + index += 1 + return token_ids + + +@pytest.mark.unit +def test_new_user_delta_survives_history_rewrite(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user must remain complete"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "second user must remain complete" + assert loss_mask == [0] * len(token_ids) + + +@pytest.mark.unit +def test_assistant_delta_keeps_existing_append_only_behavior(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "user question"}, + {"role": "assistant", "content": "reasoninganswer"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "reasoninganswer" + assert loss_mask == [1] * len(token_ids) + + +@pytest.mark.unit +def test_accumulated_multiturn_tokens_keep_every_assistant_generation_prefix(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "second reasoningsecond answer"}, + ] + + initial_prompt = tokenizer.apply_chat_template(messages[:1], add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for end in range(2, len(messages) + 1): + include_generation_prompt = messages[end - 1]["role"] == "assistant" and end > 2 + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages[:end], + include_generation_prompt=include_generation_prompt, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 2 + assert decoded.count("") == 2 + assert "first reasoning" in decoded + assert "second reasoning" in decoded + assert "second user" in decoded + + second_prefix = decoded.rindex("") + assert loss_mask[second_prefix : second_prefix + len("")] == [0] * len("") + + +@pytest.mark.unit +def test_accumulated_six_real_user_turns_keep_every_user_and_assistant_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [{"role": "user", "content": "user 1"}] + + initial_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for turn in range(1, 7): + messages.append( + { + "role": "assistant", + "content": f"reasoning {turn}answer {turn}", + } + ) + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=turn > 1, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + if turn < 6: + messages.append({"role": "user", "content": f"user {turn + 1}"}) + delta_ids, delta_mask = get_token_delta(tokenizer, messages) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 6 + assert decoded.count("") == 6 + + for turn in range(1, 7): + user_span = f"user {turn}" + user_start = decoded.index(user_span) + assert loss_mask[user_start : user_start + len(user_span)] == [0] * len(user_span) + + assistant_span = f"reasoning {turn}answer {turn}" + assistant_start = decoded.index(assistant_span) + assert loss_mask[assistant_start : assistant_start + len(assistant_span)] == [1] * len(assistant_span) + + prefix_start = decoded.rfind("", 0, assistant_start) + assert loss_mask[prefix_start:assistant_start] == [0] * (assistant_start - prefix_start) + + +@pytest.mark.unit +def test_later_assistant_allows_bpe_merge_across_generation_prefix_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = BoundaryMergingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "\nsecond answer"}, + ] + + token_ids, loss_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=True, + ) + + expected_text = "\nsecond answer" + expected_ids = tokenizer.encode(expected_text, add_special_tokens=False) + generation_prefix_length = len(tokenizer.encode("", add_special_tokens=False)) + assert token_ids == expected_ids + assert loss_mask == [0] * generation_prefix_length + [1] * (len(expected_ids) - generation_prefix_length) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_train_data_utils.py b/tests/test_train_data_utils.py new file mode 100644 index 000000000..31dc650cb --- /dev/null +++ b/tests/test_train_data_utils.py @@ -0,0 +1,393 @@ +from argparse import Namespace + +import _cp_dist_helpers # noqa: F401 +import pytest +import torch +from _cp_dist_helpers import cp_chunk_response_tensor, free_port, init_worker_process_group, stub_megatron_in_worker + +from vime.observability.train_data_utils import ( + _build_dump_payload, + restore_context_parallel_fields_to_cpu, + save_debug_train_data, +) + +NUM_GPUS = 0 + + +def _patch_single_dp_writer(monkeypatch, mpu, *, cp_size, writer_rank): + monkeypatch.setattr(mpu, "get_context_parallel_world_size", lambda: cp_size) + monkeypatch.setattr(mpu, "get_context_parallel_rank", lambda: 0) + monkeypatch.setattr(mpu, "get_context_parallel_group", lambda: None, raising=False) + monkeypatch.setattr(mpu, "is_pipeline_last_stage", lambda **_kwargs: True, raising=False) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0, raising=False) + monkeypatch.setattr( + mpu, + "get_data_parallel_rank", + lambda with_context_parallel=False: 0, + raising=False, + ) + monkeypatch.setattr( + mpu, + "get_data_parallel_world_size", + lambda with_context_parallel=False: 1, + raising=False, + ) + monkeypatch.setattr( + mpu, + "get_data_parallel_src_rank", + lambda with_context_parallel=False: writer_rank, + raising=False, + ) + + +def _restore_context_parallel_worker(rank, world_size, master_port, result_path): + import torch.distributed as dist + import torch.distributed.nn # noqa: F401 + + stub_megatron_in_worker(cp_size=world_size, cp_rank=rank) + cp_group = init_worker_process_group(rank, world_size, master_port) + try: + from megatron.core import mpu + from vime.backends.megatron_utils.cp_utils import all_gather_with_cp + + mpu.get_context_parallel_group = lambda: cp_group + + full_log_probs = torch.tensor([-0.1, -0.2, -0.3, -0.4, -0.5, -0.6]) + local_log_probs = cp_chunk_response_tensor(full_log_probs, total_length=8, response_length=6) + restored = restore_context_parallel_fields_to_cpu( + { + "total_lengths": [8], + "response_lengths": [6], + "log_probs": [local_log_probs], + }, + all_gather_with_cp, + keep_restored=rank == 0, + ) + if rank == 0: + assert restored is not None + assert restored["log_probs"][0].device.type == "cpu" + torch.save(restored["log_probs"][0], result_path) + else: + assert restored is None + finally: + dist.destroy_process_group() + + +def _single_file_dump_worker(rank, world_size, master_port, result_path): + import torch.distributed as dist + import torch.distributed.nn # noqa: F401 + + cp_size = 2 + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size=cp_size, cp_rank=cp_rank) + init_worker_process_group(rank, world_size, master_port) + cp_groups = [dist.new_group(ranks=[0, 1]), dist.new_group(ranks=[2, 3])] + dp_group = dist.new_group(ranks=[0, 2]) + try: + from megatron.core import mpu + + mpu.get_context_parallel_group = lambda: cp_groups[dp_rank] + mpu.is_pipeline_last_stage = lambda **_kwargs: True + mpu.get_tensor_model_parallel_rank = lambda: 0 + mpu.get_data_parallel_rank = lambda with_context_parallel=False: dp_rank + mpu.get_data_parallel_world_size = lambda with_context_parallel=False: 2 + mpu.get_data_parallel_src_rank = lambda with_context_parallel=False: 0 + mpu.get_data_parallel_group_gloo = lambda with_context_parallel=False: dp_group + + full_log_probs = torch.arange(6, dtype=torch.float32) + dp_rank * 10 + local_log_probs = cp_chunk_response_tensor(full_log_probs, total_length=8, response_length=6) + # Interleave sample_index across DP ranks (dp0 -> 1, dp1 -> 0) so the + # writer must sort by sample_index to restore global rollout order. + global_index = {0: 1, 1: 0}[dp_rank] + save_debug_train_data( + Namespace(save_debug_train_data=result_path), + rollout_id=7, + rollout_data={ + "tokens": [torch.arange(8) + dp_rank * 100], + "total_lengths": [8], + "response_lengths": [6], + "sample_indices": [global_index], + "log_probs": [local_log_probs], + "micro_batch_indices": [[[0]]], + }, + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_restore_context_parallel_fields_with_real_collective(tmp_path): + import torch.multiprocessing as mp + + result_path = str(tmp_path / "restored.pt") + mp.spawn( + _restore_context_parallel_worker, + args=(2, free_port(), result_path), + nprocs=2, + join=True, + ) + + torch.testing.assert_close( + torch.load(result_path, weights_only=True), + torch.tensor([-0.1, -0.2, -0.3, -0.4, -0.5, -0.6]), + ) + + +def test_save_debug_train_data_writes_one_file_with_all_dp_shards(tmp_path): + import torch.multiprocessing as mp + + result_path = str(tmp_path / "single.pt") + mp.spawn( + _single_file_dump_worker, + args=(4, free_port(), result_path), + nprocs=4, + join=True, + ) + + assert [path.name for path in tmp_path.iterdir()] == ["single.pt"] + saved = torch.load(result_path, weights_only=True) + assert saved["format_version"] == 2 + assert saved["rollout_id"] == 7 + assert saved["rank"] == 0 + # samples are sorted by sample_index across DP shards -> global rollout order. + assert [sample["sample_index"] for sample in saved["samples"]] == [0, 1] + assert [sample["data_parallel_rank"] for sample in saved["samples"]] == [1, 0] + # sample_index 0 came from dp_rank 1 (arange + 10), sample_index 1 from dp_rank 0. + torch.testing.assert_close(saved["samples"][0]["log_probs"], torch.arange(6, dtype=torch.float32) + 10) + torch.testing.assert_close(saved["samples"][1]["log_probs"], torch.arange(6, dtype=torch.float32)) + # The parallel dp_shards key keeps the DP/mbs layout without duplicating tensors. + assert [shard["rank"] for shard in saved["dp_shards"]] == [0, 2] + assert [shard["data_parallel_rank"] for shard in saved["dp_shards"]] == [0, 1] + assert [shard["sample_indices"] for shard in saved["dp_shards"]] == [[1], [0]] + assert saved["dp_shards"][0]["micro_batch_indices"] == [[[0]]] + assert "log_probs" not in saved["dp_shards"][0] + + +def test_save_debug_train_data_restores_cp_fields_by_default(tmp_path, monkeypatch): + from megatron.core import mpu + from vime.backends.megatron_utils import cp_utils + + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 3) + _patch_single_dp_writer(monkeypatch, mpu, cp_size=2, writer_rank=3) + + calls = [] + + def gather_tensor(value, total_length, response_length): + calls.append((value.clone(), total_length, response_length)) + return torch.arange(response_length, dtype=torch.float32) + total_length + + monkeypatch.setattr(cp_utils, "all_gather_with_cp", gather_tensor) + path_template = str(tmp_path / "train_{rollout_id}_{rank}.pt") + args = Namespace(save_debug_train_data=path_template) + rollout_data = { + "tokens": [torch.tensor([10, 11, 12, 13]), torch.tensor([20, 21, 22, 23, 24])], + "total_lengths": [4, 5], + "response_lengths": [2, 3], + "sample_indices": [0, 1], + "loss_masks": [torch.tensor([1, 1]), torch.tensor([1, 0, 1])], + "log_probs": [torch.tensor([-0.1]), torch.tensor([-0.2, -0.3])], + "advantages": [torch.tensor([0.5]), torch.tensor([0.6, 0.7])], + "rewards": [1.0, 0.0], + } + + save_debug_train_data(args, rollout_id=9, rollout_data=rollout_data) + + saved = torch.load(tmp_path / "train_9_3.pt", weights_only=True) + assert set(saved) == {"format_version", "rollout_id", "rank", "samples", "dp_shards"} + assert saved["format_version"] == 2 + assert saved["rollout_id"] == 9 + assert saved["rank"] == 3 + assert [sample["sample_index"] for sample in saved["samples"]] == [0, 1] + assert [sample["data_parallel_rank"] for sample in saved["samples"]] == [0, 0] + assert [sample["log_probs"].tolist() for sample in saved["samples"]] == [[4.0, 5.0], [5.0, 6.0, 7.0]] + assert [sample["advantages"].tolist() for sample in saved["samples"]] == [[4.0, 5.0], [5.0, 6.0, 7.0]] + assert [sample["rewards"] for sample in saved["samples"]] == rollout_data["rewards"] + assert saved["dp_shards"][0]["sample_indices"] == [0, 1] + # The source rollout_data tensors are left untouched (writer works on CPU copies). + torch.testing.assert_close(rollout_data["log_probs"][0], torch.tensor([-0.1])) + torch.testing.assert_close(rollout_data["log_probs"][1], torch.tensor([-0.2, -0.3])) + assert [(total, response) for _, total, response in calls] == [(4, 2), (5, 3), (4, 2), (5, 3)] + + +def test_save_debug_train_data_without_cp_uses_same_normalized_format(tmp_path, monkeypatch): + from megatron.core import mpu + + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + _patch_single_dp_writer(monkeypatch, mpu, cp_size=1, writer_rank=0) + args = Namespace( + save_debug_train_data=str(tmp_path / "no_cp_{rollout_id}_{rank}.pt"), + ) + rollout_data = { + "tokens": [torch.tensor([10, 11, 12, 13])], + "total_lengths": [4], + "response_lengths": [2], + "sample_indices": [0], + "log_probs": [torch.tensor([-0.1, -0.2], requires_grad=True)], + "advantages": [torch.tensor([0.5, 0.6])], + } + + save_debug_train_data(args, rollout_id=2, rollout_data=rollout_data) + + saved = torch.load(tmp_path / "no_cp_2_0.pt", weights_only=True) + assert set(saved) == {"format_version", "rollout_id", "rank", "samples", "dp_shards"} + assert len(saved["samples"]) == 1 + sample = saved["samples"][0] + assert sample["sample_index"] == 0 + for key in ("log_probs", "advantages"): + assert len(sample[key]) == sample["response_lengths"] + assert sample[key].device.type == "cpu" + assert not sample[key].requires_grad + torch.testing.assert_close(sample[key], rollout_data[key][0]) + + +def _layout_shard(rank, dp_rank, sample_indices, log_probs, **extra): + return { + "rank": rank, + "data_parallel_rank": dp_rank, + "rollout_data": { + "response_lengths": [len(lp) for lp in log_probs], + "sample_indices": sample_indices, + "log_probs": log_probs, + **extra, + }, + } + + +def test_build_dump_payload_sorts_samples_and_keeps_layout(): + dp_shards = [ + _layout_shard( + 0, + 0, + [2, 0], + [torch.tensor([2.0]), torch.tensor([0.0])], + micro_batch_indices=[[0, 1]], + num_microbatches=[1], + global_batch_sizes=[1], + raw_reward=[9.0], + ), + _layout_shard( + 2, + 1, + [3, 1], + [torch.tensor([3.0]), torch.tensor([1.0])], + micro_batch_indices=[[0, 1]], + num_microbatches=[1], + global_batch_sizes=[1], + raw_reward=[9.0], + ), + ] + + payload = _build_dump_payload(dp_shards, rollout_id=5, writer_rank=0) + + assert payload["format_version"] == 2 + # samples flattened across shards and sorted by global sample_index. + assert [sample["sample_index"] for sample in payload["samples"]] == [0, 1, 2, 3] + assert [sample["log_probs"].item() for sample in payload["samples"]] == [0.0, 1.0, 2.0, 3.0] + assert [sample["data_parallel_rank"] for sample in payload["samples"]] == [0, 1, 0, 1] + # per-rank scheduling stays in the parallel dp_shards key, not in samples. + assert "micro_batch_indices" not in payload["samples"][0] + assert [shard["sample_indices"] for shard in payload["dp_shards"]] == [[2, 0], [3, 1]] + assert payload["dp_shards"][0]["micro_batch_indices"] == [[0, 1]] + assert payload["dp_shards"][0]["num_microbatches"] == [1] + # whole-batch fields are stored once at the top level. + assert payload["raw_reward"] == [9.0] + + +def test_build_dump_payload_keeps_gather_order_when_index_missing(): + dp_shards = [ + _layout_shard(0, 0, None, [torch.tensor([2.0]), torch.tensor([0.0])]), + ] + + payload = _build_dump_payload(dp_shards, rollout_id=1, writer_rank=0) + + # No sample_index available -> keep DP-gather order, no sample_index key. + assert [sample["log_probs"].item() for sample in payload["samples"]] == [2.0, 0.0] + assert "sample_index" not in payload["samples"][0] + assert payload["dp_shards"][0]["sample_indices"] is None + + +def test_build_dump_payload_uses_partition_when_sample_index_missing(): + # sample_indices all None, but partition (DP positions) restores the + # exact rollout-dump order regardless of sample.index. + dp_shards = [ + { + "rank": 0, + "data_parallel_rank": 0, + "rollout_data": { + "response_lengths": [1, 1], + "sample_indices": [None, None], + "partition": [2, 0], + "log_probs": [torch.tensor([2.0]), torch.tensor([0.0])], + }, + }, + { + "rank": 2, + "data_parallel_rank": 1, + "rollout_data": { + "response_lengths": [1, 1], + "sample_indices": [None, None], + "partition": [3, 1], + "log_probs": [torch.tensor([3.0]), torch.tensor([1.0])], + }, + }, + ] + + payload = _build_dump_payload(dp_shards, rollout_id=8, writer_rank=0) + + assert [sample["rollout_position"] for sample in payload["samples"]] == [0, 1, 2, 3] + assert [sample["log_probs"].item() for sample in payload["samples"]] == [0.0, 1.0, 2.0, 3.0] + assert [shard["partition"] for shard in payload["dp_shards"]] == [[2, 0], [3, 1]] + + +def test_log_prob_capture_reorders_by_rollout_position(): + from vime.backends.megatron_utils import loss + + loss.enable_log_prob_capture() + # Two micro-batches with interleaved global positions, mirroring how the DP + # schedule strides samples across micro-batches. + loss._maybe_capture_log_probs({"partition": [2, 0]}, [torch.tensor([2.0]), torch.tensor([0.0])]) + loss._maybe_capture_log_probs({"partition": [3, 1]}, [torch.tensor([3.0]), torch.tensor([1.0])]) + captured = loss.drain_captured_log_probs() + + assert set(captured) == {0, 1, 2, 3} + # Emulate the train_actor reorder into rollout_data's sample order. + positions = [0, 1, 2, 3] + reordered = [captured[pos] for pos in positions] + assert [tensor.item() for tensor in reordered] == [0.0, 1.0, 2.0, 3.0] + + # Draining disables capture: subsequent calls are no-ops. + assert loss.drain_captured_log_probs() == {} + loss._maybe_capture_log_probs({"partition": [9]}, [torch.tensor([9.0])]) + assert loss.drain_captured_log_probs() == {} + + +def test_log_prob_capture_first_occurrence_wins(): + from vime.backends.megatron_utils import loss + + loss.enable_log_prob_capture() + loss._maybe_capture_log_probs({"partition": [0]}, [torch.tensor([1.0])]) + # A later training step recomputes the same position; the initial value wins. + loss._maybe_capture_log_probs({"partition": [0]}, [torch.tensor([9.0])]) + captured = loss.drain_captured_log_probs() + + assert captured[0].item() == 1.0 + + +def test_restore_context_parallel_fields_rejects_misaligned_samples(): + rollout_data = { + "total_lengths": [4, 5], + "response_lengths": [2, 3], + "log_probs": [torch.tensor([-0.1])], + } + + with pytest.raises(ValueError, match="one tensor per sample"): + restore_context_parallel_fields_to_cpu( + rollout_data, + lambda value, _total, _response: value, + keep_restored=True, + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_update_weight_factory.py b/tests/test_update_weight_factory.py new file mode 100644 index 000000000..0259cd892 --- /dev/null +++ b/tests/test_update_weight_factory.py @@ -0,0 +1,70 @@ +import sys +import types +from argparse import Namespace + +import pytest + +from vime.backends.megatron_utils.update_weight import create_weight_updater + +NUM_GPUS = 0 + + +class _FakeUpdater: + def __init__(self, args, model, weights_getter, *, model_name, quantization_config): + self.args = args + self.model = model + self.weights_getter = weights_getter + self.model_name = model_name + self.quantization_config = quantization_config + self.weight_version = 0 + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("mode", "transport", "colocate", "module_name", "class_name"), + [ + pytest.param("delta", "disk", False, "update_weight_from_disk_delta", "UpdateWeightFromDiskDelta", id="delta"), + pytest.param("full", "disk", False, "update_weight_from_disk", "UpdateWeightFromDisk", id="disk"), + pytest.param("full", "nccl", True, "update_weight_from_tensor", "UpdateWeightFromTensor", id="colocated"), + pytest.param( + "full", + "nccl", + False, + "update_weight_from_distributed", + "UpdateWeightFromDistributed", + id="distributed", + ), + ], +) +def test_create_weight_updater_selects_implementation(monkeypatch, mode, transport, colocate, module_name, class_name): + full_module_name = f"vime.backends.megatron_utils.update_weight.{module_name}" + fake_module = types.ModuleType(full_module_name) + setattr(fake_module, class_name, _FakeUpdater) + monkeypatch.setitem(sys.modules, full_module_name, fake_module) + + args = Namespace( + update_weight_mode=mode, + update_weight_transport=transport, + update_weight_start_version=7, + colocate=colocate, + ) + model = [object()] + + def weights_getter(): + return {"weight": object()} + + updater = create_weight_updater( + args, + model, + weights_getter, + model_name="model", + quantization_config={"quant_method": "test"}, + ) + + assert isinstance(updater, _FakeUpdater) + assert updater.args is args + assert updater.model is model + assert updater.weights_getter is weights_getter + assert updater.model_name == "model" + assert updater.quantization_config == {"quant_method": "test"} + assert updater.weight_version == 7 diff --git a/tests/test_value_temperature.py b/tests/test_value_temperature.py index cf36e4298..f4bfb47cd 100644 --- a/tests/test_value_temperature.py +++ b/tests/test_value_temperature.py @@ -1,3 +1,4 @@ +import asyncio import sys import types from argparse import Namespace @@ -9,6 +10,27 @@ NUM_GPUS = 0 +def test_opd_teacher_uses_rollout_temperature(monkeypatch): + from vime.rollout import on_policy_distillation + + captured = {} + + async def fake_post(url, payload): + captured.update(url=url, payload=payload) + return {"prompt_logprobs": []} + + monkeypatch.setattr(on_policy_distillation, "post", fake_post) + args = Namespace( + rm_url="http://teacher:8000/inference/v1/generate", + rollout_temperature=0.7, + ) + sample = Namespace(tokens=[1, 2], multimodal_inputs=None) + + asyncio.run(on_policy_distillation.reward_func(args, sample)) + + assert captured["payload"]["sampling_params"]["temperature"] == 0.7 + + def test_get_values_does_not_apply_rollout_temperature(monkeypatch): previous_loss = sys.modules.pop("vime.backends.megatron_utils.loss", None) previous_cp_utils = sys.modules.pop("vime.backends.megatron_utils.cp_utils", None) @@ -26,7 +48,7 @@ def test_get_values_does_not_apply_rollout_temperature(monkeypatch): try: from vime.backends.megatron_utils.loss import get_values - args = Namespace(qkv_format="thd", rollout_temperature=0.5, allgather_cp=False) + args = Namespace(rollout_temperature=0.5, allgather_cp=False) logits = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]], dtype=torch.float32) tokens = [torch.tensor([10, 11, 12, 13], dtype=torch.long)] diff --git a/tests/test_vllm_config_mixed_offload.py b/tests/test_vllm_config_mixed_offload.py index deed963b2..ad7ad2ad1 100644 --- a/tests/test_vllm_config_mixed_offload.py +++ b/tests/test_vllm_config_mixed_offload.py @@ -125,7 +125,6 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 8 " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( diff --git a/tests/test_vllm_config_mixed_offload_ft.py b/tests/test_vllm_config_mixed_offload_ft.py index 8f5d0145b..d2aa6ebc4 100644 --- a/tests/test_vllm_config_mixed_offload_ft.py +++ b/tests/test_vllm_config_mixed_offload_ft.py @@ -129,7 +129,6 @@ def execute(): "--actor-num-nodes 1 " "--actor-num-gpus-per-node 8 " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( diff --git a/tests/unit/rollout/test_vllm_rollout.py b/tests/test_vllm_rollout.py similarity index 56% rename from tests/unit/rollout/test_vllm_rollout.py rename to tests/test_vllm_rollout.py index e68bbea23..7b0bd392e 100644 --- a/tests/unit/rollout/test_vllm_rollout.py +++ b/tests/test_vllm_rollout.py @@ -1,18 +1,30 @@ -"""Unit tests for ``vime.rollout.vllm_rollout`` helpers and mocked async paths.""" +"""CPU unit tests for ``vime.rollout.vllm_rollout`` helpers and mocked async paths.""" from __future__ import annotations import asyncio import base64 import io +import json +import sys from argparse import Namespace from contextlib import contextmanager +from pathlib import Path from unittest.mock import AsyncMock +_tests_root = Path(__file__).resolve().parent +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs import numpy as np import pytest +_unit_stubs.install_rollout_optional_stubs() + from vime.rollout import vllm_rollout as mod + +NUM_GPUS = 0 from vime.utils.eval_config import EvalDatasetConfig from vime.utils.types import Sample @@ -116,18 +128,30 @@ def _default_sampling_params(**overrides) -> dict: return sp -def _generate_response(token_ids: list[int] | None = None) -> dict: +def _generate_response( + token_ids: list[int] | None = None, + weight_version: str | None = None, + request_spec_decode_stats: dict[str, int] | None = None, + sampling_mask: list[list[int]] | None = None, +) -> dict: tids = token_ids or [50, 51] - return { + response = { "choices": [ { "token_ids": tids, "finish_reason": "stop", - "logprobs": {"content": [{"logprob": -0.1}, {"logprob": -0.2}]}, + "logprobs": {"content": [{"logprob": -(index + 1) / 10} for index in range(len(tids))]}, } ], "usage": {"prompt_tokens": 3, "completion_tokens": len(tids)}, } + if weight_version is not None: + response["weight_version"] = weight_version + if request_spec_decode_stats is not None: + response["request_spec_decode_stats"] = request_spec_decode_stats + if sampling_mask is not None: + response["choices"][0]["sampling_mask"] = sampling_mask + return response @pytest.fixture @@ -222,6 +246,23 @@ def test_build_inference_sampling_params_forwards_disabled_top_k(): assert sp["top_k"] == -1 +@pytest.mark.unit +def test_inference_generate_tokens_and_logprobs_aligns_partial_content(): + token_ids, log_probs = mod._inference_generate_tokens_and_logprobs( + { + "token_ids": [11, 12, 13], + "logprobs": {"content": [{"logprob": -0.1}, {"logprob": -0.2}]}, + } + ) + assert token_ids == [11, 12, 13] + assert log_probs == [-0.1, -0.2, 0.0] + + +@pytest.mark.unit +def test_inference_generate_tokens_and_logprobs_rejects_invalid_token_ids(): + assert mod._inference_generate_tokens_and_logprobs({"token_ids": [1, "2"]}) == ([], []) + + @pytest.mark.unit def test_mm_render_response_to_generate_body_flat_dict(): body = mod._mm_render_response_to_generate_body( @@ -300,13 +341,24 @@ def test_mm_render_response_empty_engine_prompts_raises(): @pytest.mark.unit def test_generate_text_path_updates_sample(patch_generate_state, monkeypatch): - post_mock = AsyncMock(return_value=_generate_response([50, 51])) + post_mock = AsyncMock( + return_value=_generate_response( + [50, 51], + weight_version="step-7", + sampling_mask=[[1, 50], [2, 3, 51]], + request_spec_decode_stats={ + "num_accepted_draft_tokens": 6, + "num_draft_tokens": 8, + "num_spec_steps": 2, + }, + ) + ) monkeypatch.setattr(mod, "post", post_mock) sample = Sample(index=0, prompt="abc") result = asyncio.run( mod.generate( - _rollout_args(), + _rollout_args(vllm_speculative_config={"method": "mtp"}), sample, _default_sampling_params(max_new_tokens=8), ) @@ -315,12 +367,93 @@ def test_generate_text_path_updates_sample(patch_generate_state, monkeypatch): assert result.tokens == [97, 98, 99, 50, 51] assert result.response_length == 2 assert result.rollout_log_probs == pytest.approx([-0.1, -0.2]) + assert result.rollout_top_p_token_ids.tolist() == [1, 50, 2, 3, 51] + assert result.rollout_top_p_token_offsets.tolist() == [0, 2, 5] + assert result.weight_versions == ["step-7"] + assert result.spec_info.spec_accept_token_num == 6 + assert result.spec_info.spec_draft_token_num == 8 + assert result.spec_info.spec_verify_ct == 2 assert result.status == Sample.Status.COMPLETED body = post_mock.await_args_list[0].args[1] assert body["token_ids"] == [97, 98, 99] assert body["sampling_params"]["max_tokens"] == 8 +@pytest.mark.unit +def test_generate_streaming_records_weight_version(patch_generate_state, monkeypatch): + from vime.rollout import vllm_streaming_rollout as streaming + + class FakeStreamResponse: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + def raise_for_status(self): + return None + + async def aiter_lines(self): + chunks = [ + { + "weight_version": "step-7", + "request_spec_decode_stats": { + "num_accepted_tokens": 6, + "num_draft_tokens": 8, + "num_verify_steps": 2, + }, + "choices": [ + { + "token_ids": [50], + "sampling_mask": [[1, 50]], + "finish_reason": None, + "logprobs": {"content": [{"logprob": -0.1}]}, + } + ], + }, + { + "weight_version": "step-7", + "choices": [ + { + "token_ids": [51], + "sampling_mask": [[2, 3, 51]], + "finish_reason": "stop", + "logprobs": {"content": [{"logprob": -0.2}]}, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 2}, + }, + ] + for chunk in chunks: + yield f"data: {json.dumps(chunk)}" + yield "data: [DONE]" + + class FakeClient: + def stream(self, *args, **kwargs): + return FakeStreamResponse() + + monkeypatch.setattr(streaming, "GenerateState", _PatchedGenerateState) + monkeypatch.setattr(streaming.http_utils, "_http_client", FakeClient()) + + result = asyncio.run( + streaming.generate_streaming( + _rollout_args(vllm_speculative_config={"method": "mtp"}), + Sample(index=0, prompt="abc"), + _default_sampling_params(max_new_tokens=8), + ) + ) + + assert result.tokens == [97, 98, 99, 50, 51] + assert result.rollout_log_probs == pytest.approx([-0.1, -0.2]) + assert result.rollout_top_p_token_ids.tolist() == [1, 50, 2, 3, 51] + assert result.rollout_top_p_token_offsets.tolist() == [0, 2, 5] + assert result.weight_versions == ["step-7"] + assert result.spec_info.spec_accept_token_num == 6 + assert result.spec_info.spec_draft_token_num == 8 + assert result.spec_info.spec_verify_ct == 2 + assert result.status == Sample.Status.COMPLETED + + @pytest.mark.unit def test_generate_consistent_hash_header(patch_generate_state, monkeypatch): post_mock = AsyncMock(return_value=_generate_response()) @@ -350,7 +483,7 @@ async def fake_post(url, payload, headers=None, **kwargs): return gen_resp monkeypatch.setattr(mod, "post", fake_post) - monkeypatch.setattr(mod, "encode_image_for_rollout_engine", lambda _img: "data:image/png;base64,xx") + monkeypatch.setattr(mod, "build_multimodal_messages", lambda *_args: [{"role": "user", "content": []}]) sample = Sample(index=0, prompt="look", multimodal_inputs={"images": ["img.png"]}) result = asyncio.run(mod.generate(_rollout_args(), sample, _default_sampling_params())) @@ -359,6 +492,15 @@ async def fake_post(url, payload, headers=None, **kwargs): assert result.tokens[-1] == 13 +@pytest.mark.unit +def test_build_multimodal_messages_supports_audio_and_video(): + messages = mod.build_multimodal_messages( + "describe", + {"audio": ["https://example.com/audio.wav"], "videos": ["https://example.com/video.mp4"]}, + ) + assert [item["type"] for item in messages[0]["content"]] == ["text", "audio_url", "video_url"] + + @pytest.mark.unit def test_generate_applies_routed_experts(patch_generate_state, monkeypatch): # Fake tokenizer yields 3 prompt ids; +2 response => 5 tokens, 4 routing rows. @@ -488,6 +630,98 @@ async def custom_generate(args, sample, sampling_params, evaluation=False): assert result.reward == 0.5 +@pytest.mark.unit +def test_generate_and_rm_rejects_batched_rm_length_mismatch_without_partial_assignment( + patch_generate_state, monkeypatch +): + from vime.rollout import rm_hub + + generated_samples: list[Sample] = [] + + async def fake_generate(args, sample, sampling_params): + sample.response = "a" + sample.response_length = 1 + sample.tokens = [1] + sample.reward = None + sample.status = Sample.Status.COMPLETED + + sibling = Sample(index=1, prompt="p1", status=Sample.Status.COMPLETED) + sibling.response = "b" + sibling.response_length = 1 + sibling.tokens = [2] + sibling.reward = None + generated_samples[:] = [sample, sibling] + return generated_samples + + async def short_batched_rm(args, samples, **kwargs): + assert len(samples) == 2 + return [0.25] + + monkeypatch.setattr(mod, "generate", fake_generate) + monkeypatch.setattr(rm_hub, "load_function", lambda _path: short_batched_rm) + + with pytest.raises(ValueError, match="returned 1 rewards for 2 samples"): + asyncio.run( + mod.generate_and_rm( + _rollout_args(custom_rm_path="fake.rm"), + Sample(index=0, prompt="p0"), + _default_sampling_params(), + ) + ) + + assert [sample.reward for sample in generated_samples] == [None, None] + + +@pytest.mark.unit +@pytest.mark.parametrize( + "invalid_rewards,type_name", + [ + ({"first": 0.25, "second": 0.75}, "dict"), + ("ab", "str"), + (b"ab", "bytes"), + ], +) +def test_generate_and_rm_rejects_deceptive_batched_rm_result_types_without_partial_assignment( + patch_generate_state, monkeypatch, invalid_rewards, type_name +): + from vime.rollout import rm_hub + + generated_samples: list[Sample] = [] + + async def fake_generate(args, sample, sampling_params): + sample.response = "a" + sample.response_length = 1 + sample.tokens = [1] + sample.reward = None + sample.status = Sample.Status.COMPLETED + + sibling = Sample(index=1, prompt="p1", status=Sample.Status.COMPLETED) + sibling.response = "b" + sibling.response_length = 1 + sibling.tokens = [2] + sibling.reward = None + generated_samples[:] = [sample, sibling] + return generated_samples + + async def invalid_batched_rm(args, samples, **kwargs): + assert len(samples) == 2 + return invalid_rewards + + monkeypatch.setattr(mod, "generate", fake_generate) + monkeypatch.setattr(rm_hub, "load_function", lambda _path: invalid_batched_rm) + + with pytest.raises(TypeError, match=f"returned {type_name} instead of an iterable of rewards"): + asyncio.run( + mod.generate_and_rm( + _rollout_args(custom_rm_path="fake.rm"), + Sample(index=0, prompt="p0"), + _default_sampling_params(), + ) + ) + + assert [sample.reward for sample in generated_samples] == [None, None] + + @pytest.mark.unit def test_generate_and_rm_group_assigns_session_ids(patch_generate_state, monkeypatch): async def fake_generate_and_rm(args, sample, sampling_params, evaluation=False): @@ -503,6 +737,40 @@ async def fake_generate_and_rm(args, sample, sampling_params, evaluation=False): assert result[0].session_id != result[1].session_id +@pytest.mark.unit +def test_generate_and_rm_group_rejects_batched_rm_length_mismatch_without_partial_assignment( + patch_generate_state, monkeypatch +): + from vime.rollout import rm_hub + + async def fake_generate(args, sample, sampling_params): + sample.response = "ok" + sample.response_length = 1 + sample.tokens = [sample.index or 0] + sample.reward = None + sample.status = Sample.Status.COMPLETED + return sample + + async def long_batched_rm(args, samples, **kwargs): + assert len(samples) == 2 + return [0.25, 0.75, 1.0] + + monkeypatch.setattr(mod, "generate", fake_generate) + monkeypatch.setattr(rm_hub, "load_function", lambda _path: long_batched_rm) + + group = [Sample(index=0, prompt="p0"), Sample(index=1, prompt="p1")] + with pytest.raises(ValueError, match="returned 3 rewards for 2 samples"): + asyncio.run( + mod.generate_and_rm_group( + _rollout_args(group_rm=True, custom_rm_path="fake.rm"), + group, + _default_sampling_params(), + ) + ) + + assert [sample.reward for sample in group] == [None, None] + + @pytest.mark.unit def test_eval_rollout_passk_requests_do_not_share_session_ids(patch_generate_state, monkeypatch): seen_session_ids: list[str | None] = [] @@ -520,7 +788,12 @@ async def fake_generate_and_rm(args, sample, sampling_params, evaluation=False): args = _rollout_args() dataset_cfg = EvalDatasetConfig(name="eval", path="/tmp/eval.jsonl", n_samples_per_eval_prompt=2) - cache_key = dataset_cfg.cache_key + (args.hf_checkpoint, args.apply_chat_template) + cache_key = dataset_cfg.cache_key + ( + args.hf_checkpoint, + args.apply_chat_template, + None, + None, + ) mod.EVAL_PROMPT_DATASET[cache_key] = type("DummyDataset", (), {"samples": [Sample(prompt="prompt")]})() result = asyncio.run(mod.eval_rollout_single_dataset(args, rollout_id=0, dataset_cfg=dataset_cfg)) @@ -529,3 +802,89 @@ async def fake_generate_and_rm(args, sample, sampling_params, evaluation=False): assert None not in seen_session_ids assert len(set(seen_session_ids)) == 2 assert result[dataset_cfg.name]["samples"][0].session_id != result[dataset_cfg.name]["samples"][1].session_id + + +@pytest.mark.unit +def test_abort_deletes_inflight_without_pause_resume(patch_generate_state, monkeypatch): + from vime.backends.vllm_utils import server_control + + state = _PatchedGenerateState(_rollout_args()) + monkeypatch.setattr(mod, "GenerateState", lambda args: state) + + aborted = asyncio.Event() + posted_paths: list[str] = [] + + async def fake_get(url): + return {"workers": [{"url": "http://w0:9000"}]} + + async def fake_post(url, payload, max_retries=60, headers=None): + posted_paths.append(url) + if url.endswith("/abort_requests"): + aborted.set() + return {} + + monkeypatch.setattr(mod, "get", fake_get) + # abort() drives the delete-type sweep through the server_control helper. + monkeypatch.setattr(server_control, "post", fake_post) + + sample = Sample(index=0, prompt="p") + + async def pending_group(): + # Delete-type abort makes the in-flight /generate return on its own. + await aborted.wait() + sample.status = Sample.Status.ABORTED + return [sample] + + async def run_abort(): + state.pendings = {asyncio.create_task(pending_group())} + return await asyncio.wait_for(mod.abort(_rollout_args(), rollout_id=0), timeout=5.0) + + aborted_samples = asyncio.run(run_abort()) + + # Only /abort_requests is posted -- never /pause or /resume. + assert posted_paths and all(u.endswith("/abort_requests") for u in posted_paths) + assert state.pendings == set() + # partial_rollout is off by default, so drained groups are discarded, not returned. + assert aborted_samples == [] + + +@pytest.mark.unit +def test_abort_collects_partial_samples_when_partial_rollout(patch_generate_state, monkeypatch): + from vime.backends.vllm_utils import server_control + + args = _rollout_args(partial_rollout=True) + state = _PatchedGenerateState(args) + monkeypatch.setattr(mod, "GenerateState", lambda a: state) + + aborted = asyncio.Event() + + async def fake_get(url): + return {"workers": [{"url": "http://w0:9000"}]} + + async def fake_post(url, payload, max_retries=60, headers=None): + if url.endswith("/abort_requests"): + aborted.set() + return {} + + monkeypatch.setattr(mod, "get", fake_get) + monkeypatch.setattr(server_control, "post", fake_post) + + sample = Sample(index=0, prompt="p") + sample.response = "partial" + + async def pending_group(): + await aborted.wait() + return [sample] + + async def run_abort(): + state.pendings = {asyncio.create_task(pending_group())} + return await asyncio.wait_for(mod.abort(args, rollout_id=7), timeout=5.0) + + aborted_samples = asyncio.run(run_abort()) + + assert aborted_samples == [[sample]] + assert sample.metadata["start_rollout_id"] == 7 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/unit/backends/__init__.py b/tests/unit/backends/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/backends/megatron_utils/__init__.py b/tests/unit/backends/megatron_utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/backends/megatron_utils/update_weight/__init__.py b/tests/unit/backends/megatron_utils/update_weight/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_distributed.py b/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_distributed.py deleted file mode 100644 index 825f37094..000000000 --- a/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_distributed.py +++ /dev/null @@ -1,592 +0,0 @@ -"""Unit tests for vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py.""" - -from __future__ import annotations - -import importlib -import inspect -import sys -import types -from dataclasses import dataclass, field -from unittest.mock import MagicMock - -import pytest -import torch - -MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_distributed" - - -# Modules stubbed by _install_stubs(). These are installed ONLY for the duration of this -# module's tests (inside the fixture) and restored on teardown. Installing them at import -# time (top level) left a fake ``vllm`` (with no ``.engine``) in sys.modules, which broke -# COLLECTION of sibling test modules (e.g. test_vllm_engine.py -> ModuleNotFoundError -# 'vllm.engine'). pytest imports all test modules in one process before running fixtures, -# so the stub leak must be confined to test runtime, not collection. -_STUBBED_MODULES = ( - "megatron", - "megatron.core", - "megatron.core.parallel_state", - "megatron.core.transformer", - "megatron.core.transformer.transformer_layer", - "ray", - "ray.actor", - "vime.utils.distributed_utils", - "vllm", - "vllm.distributed", - "vllm.distributed.weight_transfer", - "vllm.distributed.weight_transfer.nccl_engine", -) - - -@pytest.fixture(scope="module") -def upw(): - saved = {k: sys.modules.get(k) for k in (*_STUBBED_MODULES, MODULE_PATH)} - # Pop first so _install_stubs()'s setdefault() actually installs the stubs (hermetic), - # then drop the module-under-test so it re-imports against the stubs. - for k in _STUBBED_MODULES: - sys.modules.pop(k, None) - _install_stubs() - sys.modules.pop(MODULE_PATH, None) - try: - yield importlib.import_module(MODULE_PATH) - finally: - for k, original in saved.items(): - if original is None: - sys.modules.pop(k, None) - else: - sys.modules[k] = original - - -def _install_stubs(): - mpu_stub = MagicMock() - mpu_stub.get_data_parallel_rank.return_value = 0 - mpu_stub.get_tensor_model_parallel_rank.return_value = 0 - mpu_stub.get_tensor_model_parallel_world_size.return_value = 1 - mpu_stub.get_pipeline_model_parallel_rank.return_value = 0 - mpu_stub.get_expert_model_parallel_world_size.return_value = 1 - mpu_stub.get_expert_model_parallel_group.return_value = "ep_group" - - megatron_core = types.ModuleType("megatron.core") - megatron_core.__path__ = [] - megatron_core.mpu = mpu_stub - parallel_state_mod = types.ModuleType("megatron.core.parallel_state") - parallel_state_mod.get_tensor_model_parallel_rank = mpu_stub.get_tensor_model_parallel_rank - parallel_state_mod.get_tensor_model_parallel_world_size = mpu_stub.get_tensor_model_parallel_world_size - transformer_mod = types.ModuleType("megatron.core.transformer") - transformer_mod.__path__ = [] - transformer_layer_mod = types.ModuleType("megatron.core.transformer.transformer_layer") - transformer_layer_mod.get_transformer_layer_offset = lambda *args, **kwargs: 0 - transformer_mod.transformer_layer = transformer_layer_mod - megatron_core.parallel_state = parallel_state_mod - megatron_core.transformer = transformer_mod - megatron_mod = types.ModuleType("megatron") - megatron_mod.core = megatron_core - sys.modules.setdefault("megatron", megatron_mod) - sys.modules.setdefault("megatron.core", megatron_core) - sys.modules.setdefault("megatron.core.parallel_state", parallel_state_mod) - sys.modules.setdefault("megatron.core.transformer", transformer_mod) - sys.modules.setdefault("megatron.core.transformer.transformer_layer", transformer_layer_mod) - - ray_mod = types.ModuleType("ray") - ray_mod.get = lambda refs: refs - ray_mod.ObjectRef = object - ray_mod.actor = types.ModuleType("ray.actor") - ray_mod.actor.ActorHandle = object - ray_mod._private = types.SimpleNamespace(services=types.SimpleNamespace(get_node_ip_address=lambda: "127.0.0.1")) - sys.modules.setdefault("ray", ray_mod) - sys.modules.setdefault("ray.actor", ray_mod.actor) - - vime_utils = types.ModuleType("vime.utils.distributed_utils") - vime_utils.get_gloo_group = MagicMock(return_value="gloo") - sys.modules.setdefault("vime.utils.distributed_utils", vime_utils) - - nccl_mod = types.ModuleType("vllm.distributed.weight_transfer.nccl_engine") - - class DummyNCCLTrainerSendWeightsArgs: - def __init__(self, *, group, packed): - self.group = group - self.packed = packed - - class DummyNCCLWeightTransferEngine: - @staticmethod - def trainer_send_weights(*args, **kwargs): - return None - - @staticmethod - def trainer_init(*args, **kwargs): - return object() - - nccl_mod.NCCLTrainerSendWeightsArgs = DummyNCCLTrainerSendWeightsArgs - nccl_mod.NCCLWeightTransferEngine = DummyNCCLWeightTransferEngine - vllm_mod = types.ModuleType("vllm") - vllm_mod.__path__ = [] - distributed_mod = types.ModuleType("vllm.distributed") - distributed_mod.__path__ = [] - weight_transfer_mod = types.ModuleType("vllm.distributed.weight_transfer") - weight_transfer_mod.__path__ = [] - vllm_mod.distributed = distributed_mod - distributed_mod.weight_transfer = weight_transfer_mod - weight_transfer_mod.nccl_engine = nccl_mod - sys.modules.setdefault("vllm", vllm_mod) - sys.modules.setdefault("vllm.distributed", distributed_mod) - sys.modules.setdefault("vllm.distributed.weight_transfer", weight_transfer_mod) - sys.modules.setdefault("vllm.distributed.weight_transfer.nccl_engine", nccl_mod) - - -@dataclass -class _RemoteCall: - args: tuple - kwargs: dict - - -class RecordingRemoteMethod: - def __init__(self, return_value: str = "ref"): - self._return_value = return_value - self.calls: list[_RemoteCall] = [] - - def remote(self, *args, **kwargs): - self.calls.append(_RemoteCall(args=args, kwargs=kwargs)) - return self._return_value - - -@dataclass -class RecordingEngine: - update_weights_from_distributed: RecordingRemoteMethod = field( - default_factory=lambda: RecordingRemoteMethod("ref") - ) - init_weights_update_group: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("init_ref")) - destroy_weights_update_group: RecordingRemoteMethod = field( - default_factory=lambda: RecordingRemoteMethod("destroy_ref") - ) - start_weight_update: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("start_ref")) - finish_weight_update: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("finish_ref")) - - -@dataclass -class RecordingLock: - acquire: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("acquired")) - release: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("released")) - - -@dataclass -class DummyGroup: - token: str = "dummy" - - -def _real_tensors(n: int = 2): - return [(f"layer.{i}.weight", torch.zeros(2, 2)) for i in range(n)] - - -def _make_dummy_nccl_engine(*, send_seen: list[dict] | None = None, init_seen: list[dict] | None = None): - """Build dummy NCCL types; patch on *upw* module (top-level import, not sys.modules).""" - - class DummyNCCLTrainerSendWeightsArgs: - def __init__(self, *, group, packed): - self.group = group - self.packed = packed - - class DummyNCCLWeightTransferEngine: - @staticmethod - def trainer_send_weights(iterator, trainer_args): - if send_seen is not None: - send_seen.append( - { - "items": list(iterator), - "group": trainer_args.group, - "packed": trainer_args.packed, - } - ) - - @staticmethod - def trainer_init(cfg): - if init_seen is not None: - init_seen.append(cfg) - return DummyGroup("group-from-trainer-init") - - return DummyNCCLWeightTransferEngine, DummyNCCLTrainerSendWeightsArgs - - -def _patch_nccl_on_module( - monkeypatch, upw, *, send_seen: list[dict] | None = None, init_seen: list[dict] | None = None -): - dummy_engine, dummy_args = _make_dummy_nccl_engine(send_seen=send_seen, init_seen=init_seen) - monkeypatch.setattr(upw, "NCCLWeightTransferEngine", dummy_engine) - monkeypatch.setattr(upw, "NCCLTrainerSendWeightsArgs", dummy_args) - - -def _patch_trainer_send(monkeypatch, upw, seen: list[dict]) -> None: - _patch_nccl_on_module(monkeypatch, upw, send_seen=seen) - monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) - - -def _make_instance(upw): - obj = object.__new__(upw.UpdateWeightFromDistributed) - obj.args = type("Args", (), {"update_weight_buffer_size": 1 << 30, "vllm_weight_sync_packed": True})() - obj.model = [] - obj.weights_getter = lambda: {} - obj.model_name = "test" - obj.quantization_config = None - obj.weight_version = 0 - obj._model_update_groups = DummyGroup() - obj._hf_weight_iterator = None - obj._is_pp_src_rank = True - obj._group_name = "g" - obj.rollout_engines = [] - obj.rollout_engine_lock = RecordingLock() - return obj - - -@pytest.mark.unit -def test_signature_no_use_vllm(upw): - sig = inspect.signature(upw.update_weights_from_distributed) - params = sig.parameters - assert "use_vllm" not in params - for p in ("group_name", "group", "weight_version", "rollout_engines", "converted_named_tensors", "packed"): - assert p in params - - -@pytest.mark.unit -def test_signature_rejects_legacy_use_vllm_call(upw): - with pytest.raises(TypeError, match="use_vllm"): - upw.update_weights_from_distributed( - "g", - DummyGroup(), - 1, - [RecordingEngine()], - _real_tensors(), - use_vllm=True, - packed=False, - ) - - -@pytest.mark.unit -def test_packed_true_uses_vllm_trainer_send_weights(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - tensors = _real_tensors() - seen = [] - _patch_trainer_send(monkeypatch, upw, seen) - - refs = upw.update_weights_from_distributed("groupA", group, 7, [engine], tensors, packed=True) - - assert len(seen) == 1 - sent = seen[0]["items"] - assert [n for n, _ in sent] == [n for n, _ in tensors] - assert seen[0]["group"] is group - assert seen[0]["packed"] is True - assert refs == ["ref"] - - -@pytest.mark.unit -def test_packed_false_still_uses_vllm_trainer_send_weights(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - tensors = _real_tensors() - seen = [] - _patch_trainer_send(monkeypatch, upw, seen) - - refs = upw.update_weights_from_distributed("groupB", group, 7, [engine], tensors, packed=False) - - assert len(seen) == 1 - assert len(seen[0]["items"]) == len(tensors) - assert seen[0]["packed"] is False - assert refs == ["ref"] - - -@pytest.mark.unit -def test_default_packed_is_false(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - seen = [] - _patch_trainer_send(monkeypatch, upw, seen) - - upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors()) - - assert len(seen) == 1 - assert seen[0]["packed"] is False - - -@pytest.mark.unit -def test_no_dist_broadcast_fallback(upw, monkeypatch): - import torch.distributed as dist - - seen_broadcast = [] - seen_send = [] - - def fake_broadcast(*a, **k): - seen_broadcast.append((a, k)) - - monkeypatch.setattr(dist, "broadcast", fake_broadcast) - _patch_trainer_send(monkeypatch, upw, seen_send) - - group = DummyGroup() - engine = RecordingEngine() - upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors(), packed=False) - - assert seen_broadcast == [] - assert len(seen_send) == 1 - - -@pytest.mark.unit -def test_remote_kwargs_include_packed_true(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - tensors = _real_tensors(n=1) - seen_send = [] - _patch_trainer_send(monkeypatch, upw, seen_send) - - upw.update_weights_from_distributed("myg", group, 42, [engine], tensors, packed=True) - - assert len(seen_send) == 1 - assert seen_send[0]["packed"] is True - assert len(engine.update_weights_from_distributed.calls) == 1 - kw = engine.update_weights_from_distributed.calls[0].kwargs - assert kw["packed"] is True - assert kw["group_name"] == "myg" - assert kw["weight_version"] == "42" - assert kw["names"] == ["layer.0.weight"] - assert kw["shapes"] == [torch.Size([2, 2])] - assert kw["dtypes"] == [torch.float32] - - -@pytest.mark.unit -def test_remote_kwargs_include_packed_false(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - tensors = _real_tensors(n=2) - seen_send = [] - _patch_trainer_send(monkeypatch, upw, seen_send) - - upw.update_weights_from_distributed("g", group, 99, [engine], tensors, packed=False) - - assert len(seen_send) == 1 - assert seen_send[0]["packed"] is False - kw = engine.update_weights_from_distributed.calls[0].kwargs - assert kw["packed"] is False - assert kw["weight_version"] == "99" - assert kw["names"] == ["layer.0.weight", "layer.1.weight"] - - -@pytest.mark.unit -def test_remote_kwargs_no_use_vllm(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - seen_send = [] - _patch_trainer_send(monkeypatch, upw, seen_send) - - upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors(), packed=False) - - assert len(seen_send) == 1 - kw = engine.update_weights_from_distributed.calls[0].kwargs - assert "use_vllm" not in kw - - -@pytest.mark.unit -def test_multiple_engines_each_get_call(upw, monkeypatch): - group = DummyGroup() - engines = [RecordingEngine() for _ in range(3)] - seen_send = [] - _patch_trainer_send(monkeypatch, upw, seen_send) - - upw.update_weights_from_distributed("g", group, 1, engines, _real_tensors(), packed=True) - assert len(seen_send) == 1 - assert seen_send[0]["packed"] is True - for e in engines: - assert len(e.update_weights_from_distributed.calls) == 1 - - -@pytest.mark.unit -def test_empty_tensor_list_still_dispatches(upw, monkeypatch): - group = DummyGroup() - engine = RecordingEngine() - seen_send = [] - _patch_trainer_send(monkeypatch, upw, seen_send) - - refs = upw.update_weights_from_distributed("g", group, 1, [engine], [], packed=False) - - assert refs == ["ref"] - kw = engine.update_weights_from_distributed.calls[0].kwargs - assert kw["names"] == [] - assert kw["shapes"] == [] - assert len(seen_send) == 1 - assert seen_send[0]["items"] == [] - - -@pytest.mark.unit -def test_raw_packed_path_sends_dense_chunks_only(upw, monkeypatch): - obj = _make_instance(upw) - obj._is_pp_src_rank = True - obj._group_name = "g" - obj._hf_weight_iterator = None - obj._use_vllm_packed = lambda: True - obj._iter_non_expert_chunks = lambda: iter([[("dense.0", torch.zeros(1))], [("dense.1", torch.zeros(1))]]) - obj._iter_expert_chunks = lambda: (_ for _ in ()).throw(AssertionError("expert pass should be skipped")) - - seen: list[tuple[list[str], bool, str]] = [] - monkeypatch.setattr( - upw.UpdateWeightFromDistributed, - "_update_bucket_weights_from_distributed", - lambda self, converted_named_tensors, pbar=None, packed=False: seen.append( - ([name for name, _ in converted_named_tensors], packed, pbar) - ), - ) - monkeypatch.setattr(upw.dist, "barrier", lambda *args, **kwargs: None) - monkeypatch.setattr(upw, "get_gloo_group", lambda: "gloo") - - upw.UpdateWeightFromDistributed._send_weights(obj, pbar="pbar") - - assert seen == [(["dense.0"], True, "pbar"), (["dense.1"], True, "pbar")] - - -@pytest.mark.unit -def test_raw_nonpacked_path_runs_dense_then_expert(upw, monkeypatch): - obj = _make_instance(upw) - obj._is_pp_src_rank = True - obj._group_name = "g" - obj._hf_weight_iterator = None - obj._use_vllm_packed = lambda: False - obj._iter_non_expert_chunks = lambda: iter([[("dense.0", torch.zeros(1))], [("dense.1", torch.zeros(1))]]) - obj._iter_expert_chunks = lambda: iter([[("expert.0", torch.zeros(1))]]) - - seen: list[tuple[list[str], bool, str]] = [] - monkeypatch.setattr( - upw.UpdateWeightFromDistributed, - "_update_bucket_weights_from_distributed", - lambda self, converted_named_tensors, pbar=None, packed=False: seen.append( - ([name for name, _ in converted_named_tensors], packed, pbar) - ), - ) - barriers: list[str] = [] - monkeypatch.setattr(upw.dist, "barrier", lambda *args, **kwargs: barriers.append(kwargs.get("group"))) - monkeypatch.setattr(upw, "get_gloo_group", lambda: "gloo") - - upw.UpdateWeightFromDistributed._send_weights(obj, pbar="pbar") - - assert seen == [ - (["dense.0"], False, "pbar"), - (["dense.1"], False, "pbar"), - (["expert.0"], False, "pbar"), - ] - assert barriers == ["gloo", "gloo"] - - -@pytest.mark.unit -def test_bridge_path_forwards_packed_flag_and_listifies_chunks(upw, monkeypatch): - obj = _make_instance(upw) - obj._is_pp_src_rank = True - obj._group_name = "g" - obj.weights_getter = lambda: {"actor": torch.zeros(1)} - obj._hf_weight_iterator = MagicMock() - obj._hf_weight_iterator.get_hf_weight_chunks.return_value = iter( - ((("bridge.0", torch.zeros(1)),), (("bridge.1", torch.zeros(1)),)) - ) - - seen: list[tuple[list[str], bool, str]] = [] - monkeypatch.setattr( - upw.UpdateWeightFromDistributed, - "_update_bucket_weights_from_distributed", - lambda self, converted_named_tensors, pbar=None, packed=False: seen.append( - ([name for name, _ in converted_named_tensors], packed, pbar) - ), - ) - monkeypatch.setattr(upw.dist, "barrier", lambda *args, **kwargs: None) - monkeypatch.setattr(upw, "get_gloo_group", lambda: "gloo") - - upw.UpdateWeightFromDistributed._sync_bridge_weights_to_rollout_engines(obj, pbar="pbar", use_vllm_packed=True) - - assert seen == [ - (["bridge.0"], True, "pbar"), - (["bridge.1"], True, "pbar"), - ] - - -@pytest.mark.unit -def test_source_no_standalone_use_vllm_param(upw): - src = inspect.getsource(upw) - lines = [line.strip() for line in src.splitlines() if "use_vllm=" in line and "use_vllm_packed" not in line] - assert lines == [] - - -@pytest.mark.unit -def test_source_no_dist_broadcast_fallback(upw): - src = inspect.getsource(upw) - assert "dist.broadcast(" not in src - - -@pytest.mark.unit -def test_source_no_materialized_named_gpu_list(upw): - src = inspect.getsource(upw.update_weights_from_distributed) - assert "named_gpu = []" not in src - assert "named_gpu_iter =" in src - - -@pytest.mark.unit -def test_connect_rollout_engines_always_uses_vllm_trainer_init(upw, monkeypatch): - args = type("Args", (), {"rollout_num_gpus_per_engine": 1})() - engines = [RecordingEngine(), RecordingEngine()] - seen: list[dict] = [] - - _patch_nccl_on_module(monkeypatch, upw, init_seen=seen) - monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(upw.torch.cuda, "empty_cache", lambda: None) - monkeypatch.setattr(upw.torch.cuda, "current_device", lambda: 0) - monkeypatch.setattr(upw.ray, "get", lambda refs: refs) - monkeypatch.setattr(upw.ray._private.services, "get_node_ip_address", lambda: "127.0.0.1") - - group = upw.connect_rollout_engines_from_distributed(args, "g", engines, engine_gpu_counts=[1, 2]) - - assert isinstance(group, DummyGroup) - assert len(seen) == 1 - assert seen[0]["master_address"] == "127.0.0.1" - assert seen[0]["world_size"] == 4 # 1 + (1 + 2) - assert len(engines[0].init_weights_update_group.calls) == 1 - assert len(engines[1].init_weights_update_group.calls) == 1 - - -@pytest.mark.unit -def test_weight_update_session_calls_start_and_finish(upw, monkeypatch): - import torch.distributed as dist - - engines = [RecordingEngine(), RecordingEngine()] - ray_refs = [] - barrier_calls: list[object] = [] - - def fake_barrier(*, group=None, **kwargs): - barrier_calls.append(group) - - monkeypatch.setattr(dist, "get_rank", lambda: 0) - monkeypatch.setattr(dist, "barrier", fake_barrier) - monkeypatch.setattr(upw, "get_gloo_group", lambda: "dummy-gloo-group") - monkeypatch.setattr(upw.ray, "get", lambda refs: ray_refs.extend(refs) or refs) - - upw._begin_vllm_weight_update_session(engines) - upw._end_vllm_weight_update_session(engines) - - assert len(engines[0].start_weight_update.calls) == 1 - assert engines[0].start_weight_update.calls[0].kwargs["is_checkpoint_format"] is True - assert len(engines[1].start_weight_update.calls) == 1 - assert len(engines[0].finish_weight_update.calls) == 1 - assert len(engines[1].finish_weight_update.calls) == 1 - assert barrier_calls == ["dummy-gloo-group", "dummy-gloo-group"] - - -@pytest.mark.unit -def test_source_wraps_sync_with_weight_update_session(upw): - src = inspect.getsource(upw.UpdateWeightFromDistributed.update_weights) - assert "_begin_vllm_weight_update_session" in src - assert "_end_vllm_weight_update_session" in src - assert "_send_weights" in src - - -@pytest.mark.unit -def test_source_uses_nccl_trainer_send_weights_args(upw): - src = inspect.getsource(upw.update_weights_from_distributed) - assert "NCCLTrainerSendWeightsArgs" in src - assert "weight_transfer_compat" not in src - - -@pytest.mark.unit -def test_cuda_sync_once_after_all_buckets_not_per_bucket(upw): - send_src = inspect.getsource(upw.update_weights_from_distributed) - sync_src = inspect.getsource(upw.UpdateWeightFromDistributed.update_weights) - assert "torch.cuda.synchronize" not in send_src - assert "torch.cuda.synchronize" in sync_src diff --git a/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py b/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py deleted file mode 100644 index 143af7645..000000000 --- a/tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py +++ /dev/null @@ -1,408 +0,0 @@ -"""Unit tests for colocated vLLM IPC weight sync.""" - -from __future__ import annotations - -import importlib -import sys -import types -from argparse import Namespace -from dataclasses import dataclass, field -from unittest.mock import MagicMock, patch - -import pytest -import torch - -MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_tensor" - -_PURGE_PREFIXES = ("megatron", "megatron_adaptor", "vime.backends.megatron_utils") - - -def _collect_subtree(prefix: str) -> list[str]: - """Collect all modules in sys.modules that start with the given prefix.""" - return [k for k in sys.modules.keys() if k == prefix or k.startswith(prefix + ".")] - - -def _install_stubs(): - mpu_stub = MagicMock() - mpu_stub.get_data_parallel_rank.return_value = 0 - mpu_stub.get_tensor_model_parallel_rank.return_value = 0 - mpu_stub.get_tensor_model_parallel_world_size.return_value = 2 - mpu_stub.get_tensor_model_parallel_group.return_value = "tp_group" - mpu_stub.get_pipeline_model_parallel_rank.return_value = 0 - - megatron_core = types.ModuleType("megatron.core") - megatron_core.__path__ = [] - megatron_core.mpu = mpu_stub - megatron_mod = types.ModuleType("megatron") - megatron_mod.__path__ = [] - megatron_mod.core = megatron_core - - sys.modules.setdefault("megatron", megatron_mod) - sys.modules.setdefault("megatron.core", megatron_core) - - ray_mod = types.ModuleType("ray") - ray_mod.get = lambda refs: refs - ray_mod.ObjectRef = object - ray_mod.actor = types.ModuleType("ray.actor") - ray_mod.actor.ActorHandle = object - sys.modules.setdefault("ray", ray_mod) - sys.modules.setdefault("ray.actor", ray_mod.actor) - - import torch.distributed as _dist - - dist_stub = MagicMock() - dist_stub.get_rank.return_value = 0 - dist_stub.get_world_size.return_value = 1 - dist_stub.get_process_group_ranks.return_value = [0, 1] - dist_stub.barrier = MagicMock() - dist_stub.all_gather_object = MagicMock() - _dist.get_rank = dist_stub.get_rank - _dist.get_world_size = dist_stub.get_world_size - _dist.get_process_group_ranks = dist_stub.get_process_group_ranks - _dist.barrier = dist_stub.barrier - _dist.all_gather_object = dist_stub.all_gather_object - - vime_utils = types.ModuleType("vime.utils.distributed_utils") - vime_utils.get_gloo_group = MagicMock(return_value="gloo") - sys.modules.setdefault("vime.utils.distributed_utils", vime_utils) - - hf_iter_stub = MagicMock() - hf_iter_stub.get_hf_weight_chunks.return_value = iter([]) - - hf_base_mod = types.ModuleType("vime.backends.megatron_utils.update_weight.hf_weight_iterator_base") - hf_base_mod.HfWeightIteratorBase = MagicMock() - hf_base_mod.HfWeightIteratorBase.create.return_value = hf_iter_stub - - upw_dist_mod = types.ModuleType("vime.backends.megatron_utils.update_weight.update_weight_from_distributed") - upw_dist_mod.connect_rollout_engines_from_distributed = MagicMock(return_value="groups") - upw_dist_mod.disconnect_rollout_engines_from_distributed = MagicMock() - upw_dist_mod.post_process_weights = MagicMock() - upw_dist_mod.update_weights_from_distributed = MagicMock(return_value=[]) - - for key, mod in [ - ("vime.backends.megatron_utils.update_weight.hf_weight_iterator_base", hf_base_mod), - ("vime.backends.megatron_utils.update_weight.update_weight_from_distributed", upw_dist_mod), - ]: - sys.modules.setdefault(key, mod) - - return hf_iter_stub, upw_dist_mod - - -_HF_ITER_STUB = MagicMock() -_HF_ITER_STUB.get_hf_weight_chunks.return_value = iter([]) - -_STUBBED_MODULES = ( - "megatron", - "megatron.core", - "ray", - "ray.actor", - "vime.utils.distributed_utils", - "vime.backends.megatron_utils.update_weight.hf_weight_iterator_base", - "vime.backends.megatron_utils.update_weight.update_weight_from_distributed", -) -_DIST_ATTRS = ("get_rank", "get_world_size", "get_process_group_ranks", "barrier", "all_gather_object") - - -@pytest.fixture(scope="module") -def upw_vllm(): - import torch.distributed as _dist - - purge_keys = set() - for prefix in _PURGE_PREFIXES: - purge_keys.update(_collect_subtree(prefix)) - for k in _STUBBED_MODULES: - purge_keys.add(k) - purge_keys.add(MODULE_PATH) - - saved_mods = {k: sys.modules.get(k) for k in purge_keys} - saved_dist = {a: getattr(_dist, a, None) for a in _DIST_ATTRS} - for k in purge_keys: - sys.modules.pop(k, None) - _install_stubs() - sys.modules.pop(MODULE_PATH, None) - - with patch("vime.utils.common.is_npu", return_value=False): - try: - yield importlib.import_module(MODULE_PATH) - finally: - for k, original in saved_mods.items(): - if original is None: - sys.modules.pop(k, None) - else: - sys.modules[k] = original - for a, original in saved_dist.items(): - if original is not None: - setattr(_dist, a, original) - - -@dataclass -class _RemoteCall: - args: tuple - kwargs: dict - - -class RecordingRemoteMethod: - def __init__(self): - self.calls: list[_RemoteCall] = [] - - def remote(self, *args, **kwargs): - self.calls.append(_RemoteCall(args=args, kwargs=kwargs)) - return "ref" - - -@dataclass -class RecordingVLLMEngine: - release_memory_occupation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - resume_memory_occupation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - init_weight_transfer_engine: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - start_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - finish_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - update_weights: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - pause_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - flush_cache: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - continue_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) - - -def _default_args(**kwargs) -> Namespace: - base = dict( - actor_num_nodes=1, - actor_num_gpus_per_node=4, - rollout_num_gpus_per_engine=2, - megatron_to_hf_mode="raw", - update_weight_buffer_size=1 << 30, - ) - base.update(kwargs) - return Namespace(**base) - - -def _make_instance(upw_vllm, args=None): - obj = object.__new__(upw_vllm.UpdateWeightFromTensor) - obj.args = args or _default_args() - obj.model = [] - obj.weights_getter = lambda: {} - obj.model_name = "test" - obj.quantization_config = None - obj.weight_version = 0 - obj._hf_weight_iterator = _HF_ITER_STUB - obj.rollout_engines = [] - obj.distributed_rollout_engines = [] - obj.use_distribute = False - obj._model_update_groups = None - obj._is_distributed_src_rank = False - obj._group_name = "vime" - obj._ipc_initialized = False - return obj - - -def _chunks(n=1): - return [[(f"p.{i}", torch.zeros(2, 2)) for i in range(2)] for _ in range(n)] - - -def _run_update(obj, *, chunks=None, rank=0) -> dict[str, int]: - chunks = chunks or _chunks(1) - obj._hf_weight_iterator = MagicMock() - obj._hf_weight_iterator.get_hf_weight_chunks.return_value = iter(chunks) - - counters = {"barrier": 0, "ipc_collect": 0} - - def counting_barrier(*args, **kwargs): - counters["barrier"] += 1 - - def counting_ipc_collect(*args, **kwargs): - counters["ipc_collect"] += 1 - - with patch("torch.distributed.get_rank", return_value=rank), patch( - "torch.distributed.barrier", side_effect=counting_barrier - ), patch("torch.cuda.ipc_collect", side_effect=counting_ipc_collect): - obj.update_weights() - return counters - - -@pytest.mark.unit -def test_colocated_lifecycle_uses_native_weight_transfer_session(upw_vllm): - obj = _make_instance(upw_vllm) - engine = RecordingVLLMEngine() - obj.rollout_engines = [engine] - - with patch(f"{MODULE_PATH}._send_to_colocated_engine") as send_to_colocated: - counters = _run_update(obj, chunks=_chunks(2)) - - assert len(engine.pause_generation.calls) == 1 - assert len(engine.flush_cache.calls) == 1 - assert len(engine.release_memory_occupation.calls) == 0 - assert len(engine.resume_memory_occupation.calls) == 0 - assert len(engine.start_weight_update.calls) == 1 - assert engine.start_weight_update.calls[0].kwargs == {"is_checkpoint_format": True} - assert len(engine.finish_weight_update.calls) == 1 - assert engine.finish_weight_update.calls[0].kwargs == {} - assert len(engine.continue_generation.calls) == 1 - - assert send_to_colocated.call_count == 2 - assert counters["ipc_collect"] == 3 - assert counters["barrier"] >= 4 - - -@dataclass -class _FakeUpdateInfo: - names: list[str] - dtype_names: list[str] - shapes: list[list[int]] - ipc_handles: list[dict[str, tuple]] - packed: bool = False - - -def _install_fake_npu_ipc_modules(monkeypatch, calls: list[dict]): - root_mod = types.ModuleType("vllm_ascend") - distributed_mod = types.ModuleType("vllm_ascend.distributed") - weight_transfer_mod = types.ModuleType("vllm_ascend.distributed.weight_transfer") - ipc_mod = types.ModuleType("vllm_ascend.distributed.weight_transfer.npu_ipc_engine") - - @dataclass - class FakeNPUIPCTrainerSendWeightsArgs: - send_mode: object - packed: bool = False - - class FakeNPUIPCWeightTransferEngine: - @staticmethod - def trainer_send_weights(iterator, trainer_args): - calls.append({"items": list(iterator), "trainer_args": trainer_args}) - trainer_args.send_mode( - _FakeUpdateInfo( - names=["layer.weight"], - dtype_names=["float32"], - shapes=[[2, 2]], - ipc_handles=[{"uuid": ("handle", ())}], - ) - ) - - ipc_mod.NPUIPCTrainerSendWeightsArgs = FakeNPUIPCTrainerSendWeightsArgs - ipc_mod.NPUIPCWeightTransferEngine = FakeNPUIPCWeightTransferEngine - weight_transfer_mod.npu_ipc_engine = ipc_mod - distributed_mod.weight_transfer = weight_transfer_mod - root_mod.distributed = distributed_mod - monkeypatch.setitem(sys.modules, "vllm_ascend", root_mod) - monkeypatch.setitem(sys.modules, "vllm_ascend.distributed", distributed_mod) - monkeypatch.setitem(sys.modules, "vllm_ascend.distributed.weight_transfer", weight_transfer_mod) - monkeypatch.setitem(sys.modules, "vllm_ascend.distributed.weight_transfer.npu_ipc_engine", ipc_mod) - - -@pytest.mark.unit -def test_send_to_colocated_engine_uses_native_npu_ipc_engine(upw_vllm, monkeypatch): - engine = RecordingVLLMEngine() - calls: list[dict] = [] - _install_fake_npu_ipc_modules(monkeypatch, calls) - - tensors = [("layer.weight", torch.zeros(2, 2))] - with patch(f"{MODULE_PATH}.is_npu", return_value=True): - upw_vllm._send_to_colocated_engine(tensors, rollout_engines=[engine], weight_version=42) - - assert calls[0]["items"] == tensors - assert calls[0]["trainer_args"].packed is False - assert len(engine.update_weights.calls) == 1 - call = engine.update_weights.calls[0] - assert call.args[0]["update_info"]["names"] == ["layer.weight"] - assert call.args[0]["update_info"]["packed"] is False - assert call.kwargs == {"weight_version": "42"} - - -@pytest.mark.unit -def test_npu_worker_patch_skips_moe_transpose_during_wake_up(upw_vllm): - wake_quant_configs = [] - - class FakeWorker: - def __init__(self): - self.vllm_config = types.SimpleNamespace(quant_config=None) - self.moe_transposed = False - - def load_model(self): - pass - - def start_weight_update(self, is_checkpoint_format=True): - pass - - def update_weights(self, update_info): - pass - - def finish_weight_update(self): - pass - - def wake_up(self, tags=None): - wake_quant_configs.append(self.vllm_config.quant_config) - if self.vllm_config.quant_config is None and (tags is None or "weights" in tags): - self.moe_transposed = True - - native_update_weights = FakeWorker.update_weights - native_finish_weight_update = FakeWorker.finish_weight_update - native_wake_up = FakeWorker.wake_up - upw_vllm._VLLMHijack._patch_one_worker(FakeWorker) - - assert FakeWorker.update_weights is native_update_weights - assert FakeWorker.finish_weight_update is native_finish_weight_update - assert FakeWorker.wake_up is not native_wake_up - - worker = FakeWorker() - worker.wake_up(tags=["weights"]) - - assert wake_quant_configs[0] is not None - assert not worker.moe_transposed - assert worker.vllm_config.quant_config is None - - -@pytest.mark.unit -def test_send_hf_params_returns_only_distributed_refs(upw_vllm): - obj = _make_instance(upw_vllm) - obj.rollout_engines = [RecordingVLLMEngine()] - obj.distributed_rollout_engines = [RecordingVLLMEngine()] - obj.use_distribute = True - obj._is_distributed_src_rank = True - obj._model_update_groups = "groups" - tensors = _chunks(1)[0] - - with patch(f"{MODULE_PATH}._send_to_colocated_engine") as send_to_colocated, patch( - f"{MODULE_PATH}.update_weights_from_distributed", return_value=["distributed-ref"] - ) as send_distributed: - refs = obj._send_hf_params(tensors) - - send_to_colocated.assert_called_once_with( - tensors, - rollout_engines=obj.rollout_engines, - weight_version=obj.weight_version, - ) - send_distributed.assert_called_once() - assert refs == ["distributed-ref"] - - -@pytest.mark.unit -def test_connect_keeps_colocated_engines_and_initializes_once(upw_vllm): - engines = [RecordingVLLMEngine() for _ in range(2)] - obj = _make_instance( - upw_vllm, - args=_default_args(actor_num_gpus_per_node=4, rollout_num_gpus_per_engine=2), - ) - - with patch("torch.distributed.get_rank", return_value=0): - obj.connect_rollout_engines( - engines, - rollout_engine_lock=MagicMock(), - engine_gpu_counts=[2, 2], - engine_gpu_offsets=[0, 2], - ) - - assert obj.rollout_engines == engines - assert obj.distributed_rollout_engines == [] - assert obj.use_distribute is False - assert obj._ipc_initialized is True - assert len(engines[0].init_weight_transfer_engine.calls) == 1 - assert len(engines[1].init_weight_transfer_engine.calls) == 1 - - engines2 = [RecordingVLLMEngine() for _ in range(2)] - with patch("torch.distributed.get_rank", return_value=0): - obj.connect_rollout_engines( - engines2, - rollout_engine_lock=MagicMock(), - engine_gpu_counts=[2, 2], - engine_gpu_offsets=[0, 2], - ) - - assert len(engines2[0].init_weight_transfer_engine.calls) == 0 - assert len(engines2[1].init_weight_transfer_engine.calls) == 0 diff --git a/tests/unit/backends/vllm_utils/__init__.py b/tests/unit/backends/vllm_utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/backends/vllm_utils/conftest.py b/tests/unit/backends/vllm_utils/conftest.py deleted file mode 100644 index ec0d8dcb8..000000000 --- a/tests/unit/backends/vllm_utils/conftest.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Shared fixtures for vLLM backend unit tests.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - - -@pytest.fixture -def vllm_args() -> SimpleNamespace: - return SimpleNamespace( - rollout_external=True, - hf_checkpoint="/tmp/model", - vllm_router_ip=None, - vllm_router_port=None, - num_gpus_per_node=8, - rollout_num_gpus_per_engine=4, - colocate=False, - debug_rollout_only=False, - actor_num_gpus_per_node=4, - actor_num_nodes=1, - use_critic=False, - critic_num_gpus_per_node=0, - critic_num_nodes=0, - ) - - -@pytest.fixture -def vllm_engine(vllm_args): - from vime.backends.vllm_utils.vllm_engine import VLLMEngine - - engine = VLLMEngine(vllm_args, rank=0) - engine.server_host = "127.0.0.1" - engine.server_port = 8765 - return engine diff --git a/tests/unit/backends/vllm_utils/test_arguments.py b/tests/unit/backends/vllm_utils/test_arguments.py deleted file mode 100644 index a438c9a75..000000000 --- a/tests/unit/backends/vllm_utils/test_arguments.py +++ /dev/null @@ -1,363 +0,0 @@ -"""Unit tests for ``vime.backends.vllm_utils.arguments``.""" - -from __future__ import annotations - -import argparse -import sys -from types import SimpleNamespace - -import pytest - - -@pytest.fixture(scope="module") -def args_mod(): - from vime.backends.vllm_utils import arguments as mod # noqa: PLC0415 - - return mod - - -@pytest.mark.unit -def test_strip_unsupported_kwargs_on_312_real(args_mod): - assert sys.version_info < (3, 13) - out = args_mod._strip_unsupported_argparse_kwargs( - {"type": int, "deprecated": True, "deprecated_aliases": ["x"], "help": "h"} - ) - assert out == {"type": int, "help": "h"} - - -@pytest.mark.unit -def test_strip_unsupported_kwargs_passthrough_on_313(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "sys", SimpleNamespace(version_info=(3, 13, 0), argv=sys.argv)) - kw = {"type": int, "deprecated": True, "help": "h"} - out = args_mod._strip_unsupported_argparse_kwargs(kw) - assert out == kw - - -@pytest.mark.unit -def test_strip_unsupported_kwargs_noop_when_absent(args_mod): - kwargs = {"type": int, "help": "h", "default": 0} - out = args_mod._strip_unsupported_argparse_kwargs(kwargs) - assert out == kwargs - - -@pytest.mark.unit -def test_wrapper_prefixes_long_flag_real_parser(args_mod): - parser = argparse.ArgumentParser(add_help=False) - wrap = args_mod._make_add_argument_wrapper(parser.add_argument) - wrap("--gpu-memory-utilization", type=float, default=0.92) - parsed, _ = parser.parse_known_args(["--vllm-gpu-memory-utilization", "0.5"]) - assert parsed.vllm_gpu_memory_utilization == 0.5 - - -@pytest.mark.unit -def test_wrapper_prefixes_dest_real_parser(args_mod): - parser = argparse.ArgumentParser(add_help=False) - wrap = args_mod._make_add_argument_wrapper(parser.add_argument) - wrap("--foo", dest="foo", type=int, default=0) - parsed, _ = parser.parse_known_args(["--vllm-foo", "7"]) - assert parsed.vllm_foo == 7 - - -@pytest.mark.unit -def test_wrapper_no_double_prefix(args_mod): - parser = argparse.ArgumentParser(add_help=False) - wrap = args_mod._make_add_argument_wrapper(parser.add_argument) - wrap("--foo", dest="vllm_foo", type=int, default=0) - dests = {a.dest for a in parser._actions if a.option_strings} - assert "vllm_foo" in dests - assert "vllm_vllm_foo" not in dests - - -@pytest.mark.unit -def test_wrapper_skips_dest_listed_in_SKIPPED_DESTS(args_mod): - parser = argparse.ArgumentParser(add_help=False) - wrap = args_mod._make_add_argument_wrapper(parser.add_argument) - wrap("--tensor-parallel-size", type=int) - flags = {s for a in parser._actions for s in a.option_strings} - assert "--vllm-tensor-parallel-size" not in flags - assert "--tensor-parallel-size" not in flags - - -@pytest.mark.unit -def test_SKIPPED_DESTS_orchestrator_parallel_dims(args_mod): - """TP/multi-node dims are orchestrator-owned; PP/DP remain CLI-forwardable.""" - assert "tensor_parallel_size" in args_mod.SKIPPED_DESTS - assert "pipeline_parallel_size" not in args_mod.SKIPPED_DESTS - assert "data_parallel_size" not in args_mod.SKIPPED_DESTS - assert "nnodes" in args_mod.SKIPPED_DESTS - assert "node_rank" in args_mod.SKIPPED_DESTS - assert "master_addr" in args_mod.SKIPPED_DESTS - assert "master_port" in args_mod.SKIPPED_DESTS - assert "data_parallel_backend" in args_mod.SKIPPED_DESTS - assert "distributed_executor_backend" in args_mod.SKIPPED_DESTS - - -@pytest.mark.unit -def test_wrapper_strips_deprecated_kwargs_when_forwarding(args_mod): - parser = argparse.ArgumentParser(add_help=False) - wrap = args_mod._make_add_argument_wrapper(parser.add_argument) - wrap("--foo", type=int, default=0, deprecated="oldname") - parsed, _ = parser.parse_known_args(["--vllm-foo", "3"]) - assert parsed.vllm_foo == 3 - - -@pytest.mark.unit -def test_detect_user_provided_value_form(args_mod): - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--foo", type=int, default=0) - user, raw = args_mod._detect_user_provided_dests(parser, ["--foo", "5"]) - assert user == {"foo"} - assert raw == {"foo": "5"} - - -@pytest.mark.unit -def test_detect_user_provided_equals_form(args_mod): - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--bar", type=str, default="x") - user, raw = args_mod._detect_user_provided_dests(parser, ["--bar=hello"]) - assert user == {"bar"} - assert raw == {"bar": "hello"} - - -@pytest.mark.unit -def test_detect_user_provided_omitted(args_mod): - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--baz", type=int, default=42) - user, raw = args_mod._detect_user_provided_dests(parser, ["--other", "1"]) - assert "baz" not in user - assert "baz" not in raw - - -@pytest.mark.unit -def test_detect_user_provided_ignores_unregistered_flags(args_mod): - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--known", type=int) - user, raw = args_mod._detect_user_provided_dests(parser, ["--unknown", "v"]) - assert user == set() - assert raw == {} - - -@pytest.mark.unit -def test_detect_user_provided_multiple(args_mod): - parser = argparse.ArgumentParser(add_help=False) - parser.add_argument("--a", type=int, default=0) - parser.add_argument("--b", type=str, default="") - user, raw = args_mod._detect_user_provided_dests(parser, ["--a", "1", "--b=hello"]) - assert user == {"a", "b"} - assert raw == {"a": "1", "b": "hello"} - - -def _ns(**overrides): - base = dict( - vllm_data_parallel_size=1, - vllm_pipeline_parallel_size=1, - rollout_num_gpus_per_engine=4, - vllm_router_ip=None, - ) - base.update(overrides) - return SimpleNamespace(**base) - - -@pytest.mark.unit -def test_validate_args_pp1(args_mod): - ns = _ns() - args_mod.validate_args(ns) - assert ns.vllm_pp_size == 1 - assert ns.vllm_dp_size == 1 - # validate_args intentionally does NOT set a global ``vllm_tp_size`` anymore: a global TP - # (derived from the *global* rollout_num_gpus_per_engine) shadowed the per-engine value and - # broke heterogeneous per-group engines (the 300s "3/4 clients joined" rendezvous hang). TP is - # now derived per engine in vllm_engine._resolve_vllm_parallel_sizes (covered in - # test_vllm_engine.py::test_resolve_parallel_sizes_is_per_engine_not_global). - assert not hasattr(ns, "vllm_tp_size") - - -@pytest.mark.unit -def test_validate_args_records_pp_dp_but_no_global_tp(args_mod): - # validate_args records pp/dp on the namespace but must not precompute a global TP, even when - # pp>1 and dp>1. Per-engine TP = gpus_per_engine // (pp * dp) is resolved at launch time. - ns = _ns(vllm_pipeline_parallel_size=2, vllm_data_parallel_size=2) - args_mod.validate_args(ns) - assert ns.vllm_pp_size == 2 - assert ns.vllm_dp_size == 2 - assert not hasattr(ns, "vllm_tp_size") - - -@pytest.mark.unit -def test_validate_args_no_longer_raises_on_pp_indivisible(args_mod): - # The pp-divisibility check moved out of validate_args and into the per-engine resolver - # (vllm_engine._resolve_vllm_parallel_sizes / compute_vllm_engine_topology), so validate_args - # itself is now agnostic to it. The enforcement is covered in test_vllm_engine.py. - ns = _ns(vllm_pipeline_parallel_size=3, rollout_num_gpus_per_engine=4) - args_mod.validate_args(ns) # must not raise - assert ns.vllm_pp_size == 3 - - -@pytest.mark.unit -def test_validate_args_router_ipv6_wrapped(args_mod): - ns = _ns(vllm_router_ip="::1") - args_mod.validate_args(ns) - assert ns.vllm_router_ip == "[::1]" - - -@pytest.mark.unit -def test_validate_args_router_ipv6_already_wrapped_unchanged(args_mod): - ns = _ns(vllm_router_ip="[::1]") - args_mod.validate_args(ns) - assert ns.vllm_router_ip == "[::1]" - - -@pytest.mark.unit -def test_validate_args_router_ipv4_unchanged(args_mod): - ns = _ns(vllm_router_ip="127.0.0.1") - args_mod.validate_args(ns) - assert ns.vllm_router_ip == "127.0.0.1" - - -@pytest.mark.unit -def test_validate_args_router_none_noop(args_mod): - ns = _ns(vllm_router_ip=None) - args_mod.validate_args(ns) - assert ns.vllm_router_ip is None - - -@pytest.mark.unit -def test_add_vllm_router_arguments_registers_vllm_prefix(args_mod): - parser = argparse.ArgumentParser(add_help=False) - args_mod.add_vllm_router_arguments(parser) - flags = {s for a in parser._actions for s in a.option_strings} - assert "--vllm-router-ip" in flags - assert "--vllm-router-port" in flags - assert "--router-request-timeout-secs" in flags - - -@pytest.mark.unit -def test_add_vllm_router_arguments_dests(args_mod): - parser = argparse.ArgumentParser(add_help=False) - args_mod.add_vllm_router_arguments(parser) - dests = {a.dest for a in parser._actions if a.option_strings} - assert "vllm_router_ip" in dests - assert "vllm_router_port" in dests - assert "router_request_timeout_secs" in dests - - -@pytest.mark.unit -def test_add_vllm_router_arguments_no_unprefixed_names(args_mod): - parser = argparse.ArgumentParser(add_help=False) - args_mod.add_vllm_router_arguments(parser) - flags = {s for a in parser._actions for s in a.option_strings} - dests = {a.dest for a in parser._actions if a.option_strings} - assert "--router-ip" not in flags - assert "--router-port" not in flags - assert "router_ip" not in dests - assert "router_port" not in dests - - -@pytest.mark.unit -def test_add_vllm_router_arguments_parses_real_values(args_mod): - parser = argparse.ArgumentParser(add_help=False) - args_mod.add_vllm_router_arguments(parser) - parsed, _ = parser.parse_known_args( - ["--vllm-router-ip", "10.0.0.1", "--vllm-router-port", "8000", "--router-request-timeout-secs", "30"] - ) - assert parsed.vllm_router_ip == "10.0.0.1" - assert parsed.vllm_router_port == 8000 - assert parsed.router_request_timeout_secs == 30 - - -@pytest.mark.unit -def test_add_vllm_router_arguments_defaults_to_consistent_hash(args_mod): - parser = argparse.ArgumentParser(add_help=False) - args_mod.add_vllm_router_arguments(parser) - parsed, _ = parser.parse_known_args([]) - assert parsed.router_policy == "consistent_hash" - - -@pytest.mark.unit -def test_orchestration_dests_use_vllm_prefix(args_mod): - assert "vllm_router_ip" in args_mod._VIME_ORCHESTRATION_DESTS - assert "vllm_router_port" in args_mod._VIME_ORCHESTRATION_DESTS - assert "router_request_timeout_secs" in args_mod._VIME_ORCHESTRATION_DESTS - assert "router_ip" not in args_mod._VIME_ORCHESTRATION_DESTS - assert "router_port" not in args_mod._VIME_ORCHESTRATION_DESTS - assert "vllm_router_request_timeout_secs" not in args_mod._VIME_ORCHESTRATION_DESTS - - -def _realistic_add_vllm_arguments(parser): - parser.add_argument("--vllm-gpu-memory-utilization", dest="vllm_gpu_memory_utilization", type=float, default=0.92) - parser.add_argument("--vllm-enforce-eager", dest="vllm_enforce_eager", action="store_true", default=False) - parser.add_argument("--vllm-router-ip", dest="vllm_router_ip", type=str, default=None) - parser.add_argument("--vllm-router-port", dest="vllm_router_port", type=int, default=None) - parser.add_argument("--vllm-server-concurrency", dest="vllm_server_concurrency", type=int, default=512) - return parser - - -@pytest.mark.unit -def test_action_table_caches(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "_VLLM_CLI_ACTION_TABLE_CACHE", None) - monkeypatch.setattr(args_mod, "add_vllm_arguments", _realistic_add_vllm_arguments) - t1 = args_mod.get_vllm_cli_action_table() - t2 = args_mod.get_vllm_cli_action_table() - assert t1 is t2 - - -@pytest.mark.unit -def test_action_table_excludes_orchestration(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "_VLLM_CLI_ACTION_TABLE_CACHE", None) - monkeypatch.setattr(args_mod, "add_vllm_arguments", _realistic_add_vllm_arguments) - table = args_mod.get_vllm_cli_action_table() - assert "vllm_gpu_memory_utilization" in table - assert "vllm_enforce_eager" in table - assert "vllm_router_ip" not in table - assert "vllm_router_port" not in table - assert "vllm_server_concurrency" not in table - - -@pytest.mark.unit -def test_action_table_flag_format(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "_VLLM_CLI_ACTION_TABLE_CACHE", None) - monkeypatch.setattr(args_mod, "add_vllm_arguments", _realistic_add_vllm_arguments) - table = args_mod.get_vllm_cli_action_table() - flag, _action = table["vllm_gpu_memory_utilization"] - assert flag == "--gpu-memory-utilization" - - -@pytest.mark.unit -def test_parse_args_tp_default_no_pp(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) - monkeypatch.setattr(sys, "argv", ["train.py", "--rollout-num-gpus-per-engine", "4"]) - ns = args_mod.vllm_parse_args() - assert ns.vllm_tensor_parallel_size == 4 - - -@pytest.mark.unit -def test_parse_args_tp_default_with_pp(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) - monkeypatch.setattr( - sys, - "argv", - ["train.py", "--rollout-num-gpus-per-engine", "4", "--vllm-pipeline-parallel-size", "2"], - ) - ns = args_mod.vllm_parse_args() - assert ns.vllm_tensor_parallel_size == 2 - - -@pytest.mark.unit -def test_parse_args_records_user_provided(args_mod, monkeypatch): - def stub(parser): - parser.add_argument("--vllm-foo", dest="vllm_foo", type=int, default=0) - return parser - - monkeypatch.setattr(args_mod, "add_vllm_arguments", stub) - monkeypatch.setattr(sys, "argv", ["train.py", "--vllm-foo", "7"]) - ns = args_mod.vllm_parse_args() - assert "vllm_foo" in ns._vllm_user_provided - assert ns._vllm_raw_values["vllm_foo"] == "7" - - -@pytest.mark.unit -def test_parse_args_default_attribute_set_even_without_register(args_mod, monkeypatch): - monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) - monkeypatch.setattr(sys, "argv", ["train.py", "--rollout-num-gpus-per-engine", "8"]) - ns = args_mod.vllm_parse_args() - assert ns.vllm_tensor_parallel_size == 8 diff --git a/tests/unit/backends/vllm_utils/test_vllm_engine.py b/tests/unit/backends/vllm_utils/test_vllm_engine.py deleted file mode 100644 index 579da57ee..000000000 --- a/tests/unit/backends/vllm_utils/test_vllm_engine.py +++ /dev/null @@ -1,684 +0,0 @@ -"""Unit tests for ``vime.backends.vllm_utils.vllm_engine``.""" - -from __future__ import annotations - -import base64 -import dataclasses -import json -import pickle - -import pytest -import requests -import torch - -from vime.backends.vllm_utils import vllm_engine as mod - - -@pytest.fixture(autouse=True) -def _seed_vllm_cli_action_table_cache(): - """Skip the device-probing rebuild of the vLLM CLI action table on CPU. - - ``get_vllm_cli_action_table()`` builds its table via ``AsyncEngineArgs.add_cli_args``, - which probes the accelerator and raises ``RuntimeError: Failed to infer device type`` - on a GPU-less host. The table only drives which ``vllm_*`` values are forwarded to the - ``vllm serve`` subprocess; the cmd/topology/sleep-mode assertions in this module exercise - vime's own explicit flag logic, not forwarding. Seed an empty (already-built) cache so the - rebuild is skipped, and restore the original on teardown so we never leak across modules. - """ - import vime.backends.vllm_utils.arguments as _args - - saved = _args._VLLM_CLI_ACTION_TABLE_CACHE - _args._VLLM_CLI_ACTION_TABLE_CACHE = {} - try: - yield - finally: - _args._VLLM_CLI_ACTION_TABLE_CACHE = saved - - -class _MockResponse: - def __init__(self, *, json_data: dict | None = None, text: str = "", status_code: int = 200): - self._json_data = json_data - self.text = text - self.status_code = status_code - # Model requests.Response.content (raw body bytes) so _response_json's empty-body - # handling (empty 200 -> {"ok": True}) is actually exercised. A JSON body is non-empty; - # text-only/empty bodies use the given text (b"" when empty). - self.content = json.dumps(json_data).encode() if json_data is not None else text.encode() - - def raise_for_status(self) -> None: - if self.status_code >= 400: - error = requests.exceptions.HTTPError(f"HTTP {self.status_code}") - error.response = self # type: ignore[assignment] - raise error - - def json(self) -> dict: - if self._json_data is None: - raise ValueError("no json") - return self._json_data - - -@pytest.mark.unit -def test_normalize_vllm_wake_tags_drops_unsupported(): - assert mod._normalize_vllm_wake_tags(["weights", "cuda_graph", "kv_cache"]) == ["weights", "kv_cache"] - - -@pytest.mark.unit -def test_normalize_vllm_wake_tags_empty_becomes_none(): - assert mod._normalize_vllm_wake_tags(["cuda_graph"]) is None - - -@pytest.mark.unit -def test_format_v6_uri_wraps_ipv6(): - assert mod._format_v6_uri("2001:db8::1") == "[2001:db8::1]" - - -@pytest.mark.unit -def test_format_v6_uri_ipv4_unchanged(): - assert mod._format_v6_uri("10.0.0.1") == "10.0.0.1" - - -@pytest.mark.unit -def test_compute_vllm_engine_topology_single_node(vllm_args): - vllm_args.num_gpus_per_node = 8 - vllm_args.rollout_num_gpus_per_engine = 4 - vllm_args.vllm_pipeline_parallel_size = 1 - vllm_args.vllm_tp_size = 4 - topo = mod.compute_vllm_engine_topology(vllm_args, global_rank=0) - assert topo.nnodes == 1 - assert topo.node_rank == 0 - assert topo.local_num_gpus == 4 - assert not topo.multi_node - assert not topo.headless - - -@pytest.mark.unit -def test_compute_vllm_engine_topology_multi_node_ranks(vllm_args): - vllm_args.num_gpus_per_node = 8 - vllm_args.rollout_num_gpus_per_engine = 16 - vllm_args.vllm_pipeline_parallel_size = 2 - vllm_args.vllm_tp_size = 8 - topo0 = mod.compute_vllm_engine_topology(vllm_args, global_rank=0) - topo1 = mod.compute_vllm_engine_topology(vllm_args, global_rank=1) - assert topo0.nnodes == 2 - assert topo0.node_rank == 0 - assert topo1.node_rank == 1 - assert topo0.local_num_gpus == 8 - assert topo0.headless is False - assert topo1.headless is True - - -@pytest.mark.unit -def test_append_distributed_flags_only_when_multi_node(vllm_args): - cmd: list[str] = ["vllm", "serve"] - single = mod.VllmEngineTopology( - nnodes=1, - node_rank=0, - local_num_gpus=4, - tensor_parallel_size=4, - pipeline_parallel_size=1, - ) - mod.append_vllm_distributed_launch_flags(cmd, single, ("10.0.0.1", 15000), vllm_args) - assert cmd == ["vllm", "serve"] - - cmd_multi: list[str] = ["vllm", "serve"] - multi = mod.VllmEngineTopology( - nnodes=2, - node_rank=1, - local_num_gpus=8, - tensor_parallel_size=8, - pipeline_parallel_size=2, - ) - mod.append_vllm_distributed_launch_flags(cmd_multi, multi, ("10.0.0.2", 16000), vllm_args) - assert "--nnodes" in cmd_multi - assert "--node-rank" in cmd_multi - assert "1" in cmd_multi - assert "--headless" in cmd_multi - assert "--master-addr" in cmd_multi - assert "10.0.0.2" in cmd_multi - assert cmd_multi[cmd_multi.index("--data-parallel-backend") + 1] == "mp" - assert cmd_multi[cmd_multi.index("--distributed-executor-backend") + 1] == "mp" - - -@pytest.mark.unit -def test_parse_dist_init_addr_ipv6(): - host, port = mod.parse_dist_init_addr("[2001:db8::1]:15000") - assert host == "2001:db8::1" - assert port == 15000 - - -@pytest.mark.unit -def test_build_vllm_subprocess_env_colocate(vllm_args, monkeypatch): - vllm_args.colocate = True - monkeypatch.delenv("PYTHONPATH", raising=False) - env = mod.build_vllm_subprocess_env( - { - "args": vllm_args, - "visible_devices": "0,1", - } - ) - assert "VLLM_ALLOW_INSECURE_SERIALIZATION" in env - assert env["VLLM_ALLOW_INSECURE_SERIALIZATION"] == "1" - assert "PYTHONPATH" in env - - -@pytest.mark.unit -def test_build_vllm_subprocess_env_sets_batch_invariant_when_deterministic(vllm_args, monkeypatch): - monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False) - vllm_args.vllm_enable_deterministic_inference = True - env = mod.build_vllm_subprocess_env({"args": vllm_args, "visible_devices": "0"}) - assert env["VLLM_BATCH_INVARIANT"] == "1" - - -@pytest.mark.unit -def test_build_vllm_subprocess_env_no_batch_invariant_by_default(vllm_args, monkeypatch): - monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False) - vllm_args.vllm_enable_deterministic_inference = False - env = mod.build_vllm_subprocess_env({"args": vllm_args, "visible_devices": "0"}) - assert "VLLM_BATCH_INVARIANT" not in env - - -@pytest.mark.unit -def test_build_vllm_cmd_adds_sleep_mode_only_for_offload_rollout(vllm_args): - vllm_args.offload_rollout = True - server_args = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) - - cmd, _ = mod.build_vllm_cmd_and_env(server_args) - - assert "--enable-sleep-mode" in cmd - assert vllm_args.vllm_enable_sleep_mode is True - - -@pytest.mark.unit -def test_build_vllm_cmd_does_not_infer_sleep_mode_from_colocate(vllm_args): - vllm_args.colocate = True - vllm_args.offload_rollout = False - server_args = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) - - cmd, _ = mod.build_vllm_cmd_and_env(server_args) - - assert "--enable-sleep-mode" not in cmd - assert not getattr(vllm_args, "vllm_enable_sleep_mode", False) - - -@pytest.mark.unit -def test_get_base_gpu_id_colocate(vllm_args): - vllm_args.colocate = True - vllm_args.num_gpus_per_node = 8 - vllm_args.rollout_num_gpus_per_engine = 4 - assert mod.get_base_gpu_id(vllm_args, rank=1) == 4 - - -@pytest.mark.unit -def test_start_weight_update_posts_four_phase_endpoint(vllm_engine, monkeypatch): - calls: list[tuple] = [] - - def fake_post(endpoint: str, payload: dict): - calls.append((endpoint, payload)) - return {"ok": True} - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post) - - result = vllm_engine.start_weight_update(is_checkpoint_format=True) - - assert result == {"ok": True} - assert len(calls) == 1 - assert calls[0][0] == "start_weight_update" - assert calls[0][1] == {"is_checkpoint_format": True} - - -@pytest.mark.unit -def test_finish_weight_update_posts_empty_body(vllm_engine, monkeypatch): - calls: list[tuple] = [] - - def fake_post(endpoint: str, payload: dict): - calls.append((endpoint, payload)) - return {"done": True} - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post) - - result = vllm_engine.finish_weight_update() - - assert result == {"done": True} - assert calls == [("finish_weight_update", {})] - - -@pytest.mark.unit -def test_update_weights_posts_native_ipc_payload_and_records_version(vllm_engine, monkeypatch): - posted: list[tuple[str, dict]] = [] - - def fake_post(endpoint: str, payload: dict): - posted.append((endpoint, payload)) - return {"ok": True} - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post) - assert vllm_engine._weight_version is None - - ipc_handles = [{"uuid-gpu0": ("rebuild_fn", (1, 2, 3))}] - vllm_engine.update_weights( - { - "update_info": { - "names": ["layer.0.weight"], - "dtype_names": ["float32"], - "shapes": [[2, 2]], - "ipc_handles": ipc_handles, - "packed": False, - } - }, - weight_version="42", - ) - - assert posted[0][0] == "update_weights" - sent = posted[0][1]["update_info"] - # ipc_handles are pickled for native vLLM/vLLM-Ascend parse_update_info. - assert "ipc_handles" not in sent - assert isinstance(sent["ipc_handles_pickled"], str) - assert pickle.loads(base64.b64decode(sent["ipc_handles_pickled"])) == ipc_handles - assert sent["names"] == ["layer.0.weight"] - assert sent["shapes"] == [[2, 2]] - assert sent["packed"] is False - # version recorded after POST success - assert vllm_engine._weight_version == "42" - - -@pytest.mark.unit -def test_update_weights_does_not_advance_version_on_failure(vllm_engine, monkeypatch): - """POST failure must not advance _weight_version (else a retry would skip the resync).""" - - def fake_post_fail(endpoint: str, payload: dict) -> dict: - raise RuntimeError("simulated POST failure") - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post_fail) - - vllm_engine._weight_version = "old" - with pytest.raises(RuntimeError, match="simulated POST failure"): - vllm_engine.update_weights( - {"update_info": {"names": [], "dtype_names": [], "shapes": [], "ipc_handles": []}}, - weight_version="new", - ) - assert vllm_engine._weight_version == "old" - - -@pytest.mark.unit -def test_get_weight_version_returns_recorded_version(vllm_engine): - vllm_engine._weight_version = "7" - assert vllm_engine.get_weight_version() == "7" - - -@pytest.mark.unit -def test_get_weight_version_raises_when_unset(vllm_engine): - """Unrecorded version is a hard error — no silent /v1/models fallback.""" - assert vllm_engine._weight_version is None - with pytest.raises(RuntimeError, match="before any successful weight transfer"): - vllm_engine.get_weight_version() - - -@pytest.mark.unit -def test_get_weight_version_worker_rank_returns_none_without_raise(vllm_engine): - """Worker ranks short-circuit (matches the class-wide idiom).""" - vllm_engine.node_rank = 1 - vllm_engine._weight_version = None - assert vllm_engine.get_weight_version() is None - - -@pytest.mark.unit -def test_update_weights_from_distributed_posts_update_weights_without_checkpoint_flag(vllm_engine, monkeypatch): - calls: list[dict] = [] - - def fake_post_vllm(update_info: dict) -> dict: - calls.append(update_info) - return {"ok": True} - - monkeypatch.setattr(vllm_engine, "_post_vllm_update_weights_http", fake_post_vllm) - - names = ["layer.0.weight"] - dtypes = [torch.float32] - shapes = [torch.Size([2, 2])] - - vllm_engine.update_weights_from_distributed( - names, - dtypes, - shapes, - group_name="vime-pp_0", - weight_version="7", - packed=True, - ) - - assert len(calls) == 1 - info = calls[0] - assert info["names"] == names - assert info["dtype_names"] == ["float32"] - assert info["shapes"] == [[2, 2]] - assert info["packed"] is True - assert "is_checkpoint_format" not in info - assert vllm_engine._weight_version == "7" - - -@pytest.mark.unit -def test_post_vllm_update_weights_http_wraps_update_info(vllm_engine, monkeypatch): - seen: list[tuple] = [] - - def fake_post(endpoint: str, payload: dict): - seen.append((endpoint, payload)) - return {"status": "ok"} - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post) - - result = vllm_engine._post_vllm_update_weights_http({"names": ["w"], "packed": False}) - - assert result == {"status": "ok"} - assert seen[0][0] == "update_weights" - assert seen[0][1] == {"update_info": {"names": ["w"], "packed": False}} - - -@pytest.mark.unit -def test_response_json_parses_dict(): - response = _MockResponse(json_data={"status": "ready"}) - assert mod._response_json(response) == {"status": "ready"} - - -@pytest.mark.unit -def test_response_json_empty_body_returns_ok(): - response = _MockResponse(text="") - assert mod._response_json(response) == {"ok": True} - - -@pytest.mark.unit -def test_response_json_invalid_json_raises(): - response = _MockResponse(text="not-json") - response.json = lambda: (_ for _ in ()).throw(ValueError("no json")) # type: ignore[method-assign] - with pytest.raises(ValueError, match="no json"): - mod._response_json(response) - - -@pytest.mark.unit -def test_response_json_http_error_adds_response_text_note(): - response = _MockResponse(json_data={"error": "bad"}, text="server exploded", status_code=500) - with pytest.raises(requests.exceptions.HTTPError) as exc_info: - mod._response_json(response) - assert "response.text='server exploded'" in exc_info.value.__notes__ - - -@pytest.mark.unit -def test_http_base_ipv6_host(vllm_engine): - vllm_engine.server_host = "[2001:db8::1]" - assert vllm_engine._http_base() == "http://[2001:db8::1]:8765" - - -@pytest.mark.unit -def test_redact_cmd_for_log_masks_hf_token(): - cmd = ["vllm", "serve", "model", "--hf-token", "secret-token", "--port", "8000"] - logged = mod.redact_cmd_for_log(cmd) - assert "secret-token" not in logged - assert "***" in logged - - -@pytest.mark.unit -def test_serialize_for_cli_primitives(): - assert mod.serialize_for_cli(42) == "42" - assert mod.serialize_for_cli(True) == "True" - assert mod.serialize_for_cli({"backend": "nccl"}) == json.dumps({"backend": "nccl"}) - - -@pytest.mark.unit -def test_serialize_for_cli_dataclass(): - @dataclasses.dataclass - class _Cfg: - backend: str = "nccl" - - out = mod.serialize_for_cli(_Cfg()) - assert json.loads(out) == {"backend": "nccl"} - - -@pytest.mark.unit -def test_get_base_gpu_id_with_critic_offset(vllm_args): - vllm_args.colocate = False - vllm_args.debug_rollout_only = False - vllm_args.actor_num_gpus_per_node = 4 - vllm_args.actor_num_nodes = 1 - vllm_args.use_critic = True - vllm_args.critic_num_gpus_per_node = 2 - vllm_args.critic_num_nodes = 1 - vllm_args.num_gpus_per_node = 8 - vllm_args.rollout_num_gpus_per_engine = 2 - # actor 4 + critic 2 + rank0*2 = 6 - assert mod.get_base_gpu_id(vllm_args, rank=0) == 6 - - -@pytest.mark.unit -def test_resume_memory_occupation_wake_tags_query(vllm_engine, monkeypatch): - seen: list[tuple] = [] - - def fake_post(url, *, params=None, timeout=30, json=None): - seen.append((url, params, timeout, json)) - return _MockResponse(json_data={"ok": True}) - - vllm_args = vllm_engine.args - vllm_args.vllm_enable_sleep_mode = True - monkeypatch.setattr(mod.requests, "post", fake_post) - - vllm_engine.resume_memory_occupation(tags=["weights", "cuda_graph"]) - - assert len(seen) == 1 - assert seen[0][1] == [("tags", "weights")] - - -@pytest.mark.unit -def test_release_memory_occupation_flushes_then_posts_sleep(vllm_engine, monkeypatch): - calls: list[str] = [] - - def fake_flush_cache(): - calls.append("flush_cache") - - def fake_post(url, *, params=None, timeout=30, json=None): - calls.append(url) - assert params == {"level": 2} - assert timeout == 30 - assert json is None - return _MockResponse(json_data={"ok": True, "sleep_mode": True}) - - vllm_engine.args.vllm_enable_sleep_mode = False - monkeypatch.setattr(vllm_engine, "flush_cache", fake_flush_cache) - monkeypatch.setattr(mod.requests, "post", fake_post) - - assert vllm_engine.release_memory_occupation(level=2) == {"ok": True, "sleep_mode": True} - assert calls == ["flush_cache", "http://127.0.0.1:8765/sleep"] - - -@pytest.mark.unit -def test_resume_memory_occupation_posts_wake_even_when_sleep_disabled(vllm_engine, monkeypatch): - seen: list[tuple] = [] - - def fake_post(url, *, params=None, timeout=30, json=None): - seen.append((url, params, timeout, json)) - return _MockResponse(json_data={"ok": True, "sleep_mode": True}) - - vllm_engine.args.vllm_enable_sleep_mode = False - monkeypatch.setattr(mod.requests, "post", fake_post) - - assert vllm_engine.resume_memory_occupation() == {"ok": True, "sleep_mode": True} - assert seen == [("http://127.0.0.1:8765/wake_up", None, 30, None)] - - -@pytest.mark.unit -def test_init_weights_update_group_retries_then_succeeds(vllm_engine, monkeypatch): - attempts = {"n": 0} - - def fake_post(endpoint: str, payload: dict): - attempts["n"] += 1 - if attempts["n"] < 2: - raise requests.ConnectionError("transient") - return {"initialized": True} - - monkeypatch.setattr(vllm_engine, "_make_request", fake_post) - monkeypatch.setattr(mod.time, "sleep", lambda *_a, **_k: None) - - result = vllm_engine.init_weights_update_group( - "127.0.0.1", - 29500, - rank_offset=1, - world_size=4, - group_name="unused", - backend="nccl", - ) - - assert result == {"initialized": True} - assert attempts["n"] == 2 - - -@pytest.mark.unit -def test_init_weights_update_group_raises_after_three_failures(vllm_engine, monkeypatch): - monkeypatch.setattr( - vllm_engine, - "_make_request", - lambda *a, **k: (_ for _ in ()).throw(requests.ConnectionError("down")), - ) - monkeypatch.setattr(mod.time, "sleep", lambda *_a, **_k: None) - - with pytest.raises(RuntimeError, match="init_weight_transfer_engine failed"): - vllm_engine.init_weights_update_group( - "127.0.0.1", - 29500, - rank_offset=1, - world_size=4, - group_name="g", - backend="nccl", - ) - - -def _stub_server_info(monkeypatch, parallel_config: dict) -> None: - def fake_get(url, *, params=None, timeout=30): - return _MockResponse(json_data={"vllm_config": {"parallel_config": parallel_config}}) - - monkeypatch.setattr(mod.requests, "get", fake_get) - - -@pytest.mark.unit -def test_sanity_check_external_server_args_passes_on_match(vllm_engine, monkeypatch): - vllm_engine._server_args = {"tp_size": 2, "pp_size": 1, "dp_size": 1, "nnodes": 1} - _stub_server_info( - monkeypatch, - {"tensor_parallel_size": 2, "pipeline_parallel_size": 1, "data_parallel_size": 1, "nnodes": 1}, - ) - # All reported fields match the per-engine expectation → no raise. - vllm_engine._sanity_check_external_server_args() - - -@pytest.mark.unit -def test_sanity_check_external_server_args_raises_on_tp_mismatch(vllm_engine, monkeypatch): - # Expect a tp=2 engine but the external server reports tp=1 → fail fast (the bug class - # that used to only warn and then hang the weight-sync rendezvous 300s later). - vllm_engine._server_args = {"tp_size": 2, "pp_size": 1, "dp_size": 1, "nnodes": 1} - _stub_server_info( - monkeypatch, - {"tensor_parallel_size": 1, "pipeline_parallel_size": 1, "data_parallel_size": 1, "nnodes": 1}, - ) - with pytest.raises(AssertionError, match="tp_size"): - vllm_engine._sanity_check_external_server_args() - - -@pytest.mark.unit -def test_sanity_check_external_server_args_skips_unreported_field(vllm_engine, monkeypatch): - # vLLM /server_info may not surface ``nnodes``: an unreported (None) field is skipped, - # not treated as a mismatch — so a single-node external engine doesn't false-fail. - vllm_engine._server_args = {"tp_size": 1, "pp_size": 1, "dp_size": 1, "nnodes": 2} - _stub_server_info( - monkeypatch, - {"tensor_parallel_size": 1, "pipeline_parallel_size": 1, "data_parallel_size": 1}, - ) - vllm_engine._sanity_check_external_server_args() - - -@pytest.mark.unit -def test_sanity_check_external_server_args_raises_when_parallel_config_missing(vllm_engine, monkeypatch): - vllm_engine._server_args = {"tp_size": 1, "pp_size": 1, "dp_size": 1, "nnodes": 1} - _stub_server_info(monkeypatch, {}) - with pytest.raises(RuntimeError, match="missing vllm_config.parallel_config"): - vllm_engine._sanity_check_external_server_args() - - -@pytest.mark.unit -def test_resolve_parallel_sizes_is_per_engine_not_global(vllm_args): - # The global flag is 1, but THIS engine has 2 GPUs → tp must be 2 (per-engine), not 1. - # A stale global vllm_tp_size must NOT shadow the per-engine value. - vllm_args.rollout_num_gpus_per_engine = 1 - vllm_args.vllm_pipeline_parallel_size = 1 - vllm_args.vllm_tp_size = 1 # stale global; must be ignored now - tp, pp = mod._resolve_vllm_parallel_sizes(vllm_args, gpus_per_engine=2) - assert (tp, pp) == (2, 1) - - -@pytest.mark.unit -def test_compute_topology_heterogeneous_per_group_tp(vllm_args): - # Reproduces the rendezvous bug's root: global=1 but a per-group engine uses 2 GPUs. - vllm_args.num_gpus_per_node = 8 - vllm_args.rollout_num_gpus_per_engine = 1 - vllm_args.vllm_pipeline_parallel_size = 1 - vllm_args.vllm_tp_size = 1 # stale global - topo = mod.compute_vllm_engine_topology(vllm_args, global_rank=0, num_gpus_per_engine=2) - assert topo.tensor_parallel_size == 2 - assert topo.nnodes == 1 - - -@pytest.mark.unit -def test_resolve_parallel_sizes_dp_consumes_gpus(vllm_args): - # vLLM DP consumes GPUs (total = tp * pp * dp), so tp = gpus // (pp * dp). - # dp=2, pp=1, 4 GPUs/engine → tp=2. - vllm_args.vllm_pipeline_parallel_size = 1 - vllm_args.vllm_data_parallel_size = 2 - vllm_args.vllm_dp_size = 2 - tp, pp = mod._resolve_vllm_parallel_sizes(vllm_args, gpus_per_engine=4) - assert (tp, pp) == (2, 1) - - -@pytest.mark.unit -def test_resolve_parallel_sizes_dp_and_pp_combined(vllm_args): - # dp=2, pp=2, 8 GPUs/engine → tp = 8 // (2*2) = 2. - vllm_args.vllm_pipeline_parallel_size = 2 - vllm_args.vllm_data_parallel_size = 2 - vllm_args.vllm_dp_size = 2 - tp, pp = mod._resolve_vllm_parallel_sizes(vllm_args, gpus_per_engine=8) - assert (tp, pp) == (2, 2) - - -@pytest.mark.unit -def test_resolve_parallel_sizes_rejects_indivisible_dp(vllm_args): - # gpus_per_engine not divisible by pp*dp must raise (fail fast, not desync the rendezvous). - vllm_args.vllm_pipeline_parallel_size = 1 - vllm_args.vllm_data_parallel_size = 2 - vllm_args.vllm_dp_size = 2 - with pytest.raises(ValueError, match="divisible"): - mod._resolve_vllm_parallel_sizes(vllm_args, gpus_per_engine=3) - - -@pytest.mark.unit -def test_make_request_short_circuits_on_headless(vllm_engine, monkeypatch): - # _make_request is the single control-plane POST choke point; on a headless worker - # (node_rank>0) it must no-op to None without issuing any HTTP request. - def _boom(*a, **k): - raise AssertionError("control-plane HTTP must not be called on a headless worker") - - monkeypatch.setattr(mod.requests, "post", _boom) - vllm_engine.node_rank = 1 - assert vllm_engine._make_request("whatever", {}) is None - - -@pytest.mark.unit -def test_control_plane_methods_noop_on_headless_worker(vllm_engine, monkeypatch): - """node_rank>0 (headless) workers own no HTTP server; every control-plane method must - no-op (return None) without issuing an HTTP request.""" - - def _boom(*a, **k): - raise AssertionError("control-plane HTTP must not be called on a headless worker") - - monkeypatch.setattr(mod.requests, "post", _boom) - monkeypatch.setattr(mod.requests, "get", _boom) - vllm_engine.node_rank = 1 - vllm_engine.args.vllm_enable_sleep_mode = True - - assert vllm_engine.init_weight_transfer_engine({"init_info": {}}) is None - assert vllm_engine.start_weight_update() is None - assert vllm_engine.finish_weight_update() is None - assert vllm_engine.init_weights_update_group("addr", 1, 0, 4, "g", "nccl") is None - assert vllm_engine.update_weights_from_distributed(["w"], [torch.float32], [[1]], "g") is None - assert vllm_engine.release_memory_occupation() is None - assert vllm_engine.resume_memory_occupation() is None diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py deleted file mode 100644 index 0ef7825be..000000000 --- a/tests/unit/conftest.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Unit-test collection hooks (e.g. stub optional heavy deps on dev machines).""" - -from __future__ import annotations - -import importlib.util -import sys -from unittest.mock import MagicMock - - -def _ensure_ray_stub() -> None: - # Only stub ray when it is genuinely absent (bare-deps dev machine). When the real ray is - # installed (as in the CI image), a bare MagicMock stub shadows it and breaks unrelated - # submodule imports like ``from ray.util.placement_group import placement_group`` in - # sibling test packages — this conftest's module-level mutation leaks session-wide. - if "ray" in sys.modules: - return - try: - if importlib.util.find_spec("ray") is not None: - return - except (ImportError, ValueError): - pass - ray = MagicMock() - sys.modules["ray"] = ray - sys.modules["ray._private"] = MagicMock() - sys.modules["ray._private.services"] = MagicMock() - sys.modules["ray.actor"] = MagicMock() - - -_ensure_ray_stub() diff --git a/tests/unit/rollout/conftest.py b/tests/unit/rollout/conftest.py deleted file mode 100644 index aa10b668f..000000000 --- a/tests/unit/rollout/conftest.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Shared stubs for ``vime.rollout.vllm_rollout`` unit tests.""" - -from __future__ import annotations - -import importlib.util -import sys -import types -from unittest.mock import MagicMock - - -def _real_module_available(name: str) -> bool: - """True if the real module is importable, so we should NOT shadow it with a stub. - - These conftest stubs install bare ``sys.modules`` entries at import time and never - clean up, so they leak to sibling test packages (e.g. ``tests/utils``). When the real - dependency is installed (as in the CI image), a bare stub like ``vllm_router`` (no - submodules) breaks unrelated imports such as ``from vllm_router.launch_router import``. - Only stub when the real package is genuinely absent (bare-deps dev machines). - """ - if name in sys.modules: - return True - try: - return importlib.util.find_spec(name) is not None - except (ImportError, ValueError): - return False - - -def _ensure_vllm_router_stub() -> None: - if _real_module_available("vllm_router"): - return - sys.modules["vllm_router"] = types.ModuleType("vllm_router") - - -def _ensure_pil_stub() -> None: - if _real_module_available("PIL"): - return - pil = types.ModuleType("PIL") - image_mod = types.ModuleType("PIL.Image") - pil.Image = image_mod - sys.modules["PIL"] = pil - sys.modules["PIL.Image"] = image_mod - - -def _ensure_transformers_stub() -> None: - if _real_module_available("transformers"): - return - mod = types.ModuleType("transformers") - mod.AutoTokenizer = type( - "AutoTokenizer", - (), - {"from_pretrained": staticmethod(lambda *args, **kwargs: object())}, - ) - mod.AutoProcessor = type( - "AutoProcessor", - (), - {"from_pretrained": staticmethod(lambda *args, **kwargs: (_ for _ in ()).throw(OSError()))}, - ) - mod.PreTrainedTokenizerBase = type("PreTrainedTokenizerBase", (), {}) - mod.ProcessorMixin = type("ProcessorMixin", (), {}) - sys.modules["transformers"] = mod - - -def _ensure_aiohttp_stub() -> None: - if _real_module_available("aiohttp"): - return - sys.modules["aiohttp"] = MagicMock() - - -def _ensure_pylatexenc_stub() -> None: - if _real_module_available("pylatexenc"): - return - pylatexenc = types.ModuleType("pylatexenc") - latex2text = types.ModuleType("pylatexenc.latex2text") - pylatexenc.latex2text = latex2text - sys.modules["pylatexenc"] = pylatexenc - sys.modules["pylatexenc.latex2text"] = latex2text - - -_ensure_vllm_router_stub() -_ensure_pil_stub() -_ensure_transformers_stub() -_ensure_aiohttp_stub() -_ensure_pylatexenc_stub() diff --git a/tests/utils/test_hf_checkpoint_saver.py b/tests/utils/test_hf_checkpoint_saver.py index 7b8851537..9f9767d66 100644 --- a/tests/utils/test_hf_checkpoint_saver.py +++ b/tests/utils/test_hf_checkpoint_saver.py @@ -1,5 +1,6 @@ import json from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -8,10 +9,12 @@ from vime.backends.megatron_utils.hf_checkpoint_saver import ( _clear_existing_hf_weights, _copy_hf_assets, + _finalize_local_shards, _SafetensorShardWriter, + _write_pending_chunk, + save_hf_model_to_path, ) - NUM_GPUS = 0 @@ -51,10 +54,17 @@ def test_clear_existing_hf_weights_removes_old_weight_files_only(tmp_path: Path) assert not (tmp_path / "pytorch_model.bin").exists() +def test_save_hf_model_to_path_rejects_origin_checkpoint(tmp_path: Path): + args = SimpleNamespace(hf_checkpoint=str(tmp_path)) + + with pytest.raises(ValueError, match="same directory as --hf-checkpoint"): + save_hf_model_to_path(args, tmp_path, model=None) + + def test_safetensor_shard_writer_writes_hf_index(tmp_path: Path): writer = _SafetensorShardWriter(tmp_path, enabled=True) - writer.write([("layers.0.weight", torch.ones(2, 2)), ("layers.0.weight_scale", torch.ones(1))]) - writer.write([("layers.1.weight", torch.zeros(2, 2))]) + writer.write([("layers.0.weight", torch.ones(2, 2)), ("layers.0.weight_scale", torch.ones(1))], shard_idx=0) + writer.write([("layers.1.weight", torch.zeros(2, 2))], shard_idx=1) writer.finalize() index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) @@ -71,5 +81,59 @@ def test_safetensor_shard_writer_writes_hf_index(tmp_path: Path): assert torch.equal(shard1["layers.1.weight"], torch.zeros(2, 2)) +def test_finalize_shard_files_merges_node_writer_states(tmp_path: Path): + writer0 = _SafetensorShardWriter(tmp_path, enabled=True) + writer1 = _SafetensorShardWriter(tmp_path, enabled=True) + + writer0.write([("layers.0.weight", torch.ones(2, 2))], shard_idx=0) + writer1.write([("layers.1.weight", torch.zeros(2, 2))], shard_idx=1) + + # each rank renames its own files off the shared plan; rank 0 writes the index + states = [writer0.state(), writer1.state()] + for rank, state in enumerate(states): + _finalize_local_shards(tmp_path, state, states, write_index=rank == 0) + + index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) + assert index["metadata"]["total_size"] == 32 + assert index["weight_map"] == { + "layers.0.weight": "model-00001-of-00002.safetensors", + "layers.1.weight": "model-00002-of-00002.safetensors", + } + assert not (tmp_path / "model-00001.safetensors").exists() + assert not (tmp_path / "model-00002.safetensors").exists() + + shard0 = load_file(tmp_path / "model-00001-of-00002.safetensors") + shard1 = load_file(tmp_path / "model-00002-of-00002.safetensors") + assert torch.equal(shard0["layers.0.weight"], torch.ones(2, 2)) + assert torch.equal(shard1["layers.1.weight"], torch.zeros(2, 2)) + + +def test_pending_chunk_write_flushes_incomplete_node_group(tmp_path: Path): + num_nodes = 3 + writers = [_SafetensorShardWriter(tmp_path, enabled=True) for _ in range(num_nodes)] + pending_writes = [None] * num_nodes + + for chunk_idx in range(5): + node_rank = chunk_idx % num_nodes + pending_writes[node_rank] = ( + chunk_idx, + [(f"layers.{chunk_idx}.weight", torch.full((1,), chunk_idx, dtype=torch.float32))], + ) + + if (chunk_idx + 1) % num_nodes == 0: + for i, writer in enumerate(writers): + pending_writes[i] = _write_pending_chunk(writer, pending_writes[i]) + + for i, writer in enumerate(writers): + pending_writes[i] = _write_pending_chunk(writer, pending_writes[i]) + + states = [writer.state() for writer in writers] + for rank, state in enumerate(states): + _finalize_local_shards(tmp_path, state, states, write_index=rank == 0) + + index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) + assert index["weight_map"] == {f"layers.{i}.weight": f"model-{i + 1:05d}-of-00005.safetensors" for i in range(5)} + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_loss_mask_type_qwen35.py b/tests/utils/test_loss_mask_type_qwen35.py index 6aae0faca..2c487a956 100644 --- a/tests/utils/test_loss_mask_type_qwen35.py +++ b/tests/utils/test_loss_mask_type_qwen35.py @@ -1,5 +1,9 @@ +import pytest + from vime.utils.mask_utils import MultiTurnLossMaskGenerator +NUM_GPUS = 0 + class FakeQwen35Tokenizer: """A tiny char-level tokenizer that models the Qwen3.5 assistant formatting rule. @@ -234,3 +238,34 @@ def test_qwen3_5_matches_expected_mask_for_tool_call_flow(): "\n\n\nTOOL_CALL\n\n\n\n\nls\n\n\n<|im_end|>\n", "REASONING\n\n\nFINAL<|im_end|>\n", ] + + +def test_qwen3_matches_full_template_for_consecutive_tool_responses(): + tokenizer = FakeQwen35Tokenizer() + messages = [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "USER"}, + { + "role": "assistant", + "content": "CALL", + "tool_calls": [ + {"function": {"name": "terminal", "arguments": {"command": "cat a.py"}}}, + {"function": {"name": "terminal", "arguments": {"command": "cat b.py"}}}, + ], + }, + {"role": "tool", "content": "CONTENTS_A"}, + {"role": "tool", "content": "CONTENTS_B"}, + ] + + expected_text, expected_loss_mask = tokenizer.render_with_expected_mask(messages) + expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"] + generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3") + + token_ids, loss_mask = generator.get_loss_mask(messages) + + assert token_ids == expected_token_ids + assert loss_mask == expected_loss_mask + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_mask_utils.py b/tests/utils/test_mask_utils.py deleted file mode 100644 index cb5f5608b..000000000 --- a/tests/utils/test_mask_utils.py +++ /dev/null @@ -1,99 +0,0 @@ -from transformers import AutoTokenizer - -from vime.utils.mask_utils import MultiTurnLossMaskGenerator - - -def test_loss_mask_qwen3_simple(model_name: str = "Qwen/Qwen3-8B"): - tokenizer = AutoTokenizer.from_pretrained(model_name) - mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3") - messages = [ - {"role": "system", "content": "SYSTEM MESSAGE FOR TESTING ONLY"}, - {"role": "user", "content": "USER CONTENT FOR TESTING ONLY"}, - {"role": "assistant", "content": "ASSISTANT RESPONSE FOR TESTING ONLY"}, - ] - all_token_ids, all_loss_masks = mask_generator.gen_multi_turn_loss_mask_qwen3(messages) - assert len(all_token_ids) == len(all_loss_masks), f"{len(all_token_ids)} != {len(all_loss_masks)}" - selected_texts = mask_generator.get_text_from_loss_mask(all_token_ids, all_loss_masks) - assert len(selected_texts) == 1, f"Expected 1 text, got {len(selected_texts)}" - - print(f"==== Single Turn Test {model_name} ====") - print("text = ", [tokenizer.decode(all_token_ids)]) - print("token_ids = ", all_token_ids) - print("loss_mask = ", all_loss_masks) - print("selected_texts = ", selected_texts) - - -def test_loss_mask_qwen3_tools(model_name: str = "Qwen/Qwen3-8B"): - tokenizer = AutoTokenizer.from_pretrained(model_name) - mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3") - messages = [ - {"role": "system", "content": "SYSTEM MESSAGE FOR TESTING ONLY"}, - {"role": "user", "content": "USER CONTENT FOR TESTING ONLY"}, - { - "role": "assistant", - "content": "I WILL CALL terminal", - "tool_calls": [ - {"function": {"name": "terminal", "arguments": {"command": "ls"}}, "id": "call_0", "type": "function"}, - {"function": {"name": "terminal", "arguments": {"command": "ls"}}, "id": "call_0", "type": "function"}, - ], - }, - {"role": "tool", "name": "terminal", "content": "LICENSE README.md README_zh.md"}, - {"role": "tool", "name": "terminal", "content": "LICENSE README.md README_zh.md"}, - {"role": "assistant", "content": "ASSISTANT RESPONSE FOR TESTING ONLY"}, - ] - tools = [ - { - "type": "function", - "function": { - "name": "terminal", - "description": "Perform operations from the terminal.", - "parameters": { - "type": "object", - "properties": { - "command": { - "type": "string", - "description": "The bash command to execute as `bash -c `", - }, - "description": { - "type": "string", - "description": "Brief description of the command for the user.", - }, - }, - "required": ["command"], - }, - }, - }, - { - "type": "function", - "function": { - "name": "read_file", - "description": "Read the content of a file given its path.", - "parameters": { - "type": "object", - "properties": { - "file_path": { - "type": "string", - "description": "The absolute path to the file to be read.", - } - }, - "required": ["file_path"], - }, - }, - }, - ] - - all_token_ids, all_loss_masks = mask_generator.gen_multi_turn_loss_mask_qwen3(messages, tools) - assert len(all_token_ids) == len(all_loss_masks), f"{len(all_token_ids)} != {len(all_loss_masks)}" - selected_texts = mask_generator.get_text_from_loss_mask(all_token_ids, all_loss_masks) - assert len(selected_texts) == 2, f"Expected 2 texts, got {len(selected_texts)}" - - print(f"==== Multi-turn with Tools Test {model_name} ====") - print("text = ", [tokenizer.decode(all_token_ids)]) - print("token_ids = ", all_token_ids) - print("loss_mask = ", all_loss_masks) - print("selected_texts = ", selected_texts) - - -if __name__ == "__main__": - test_loss_mask_qwen3_simple("Qwen/Qwen3-Coder-30B-A3B-Instruct") - test_loss_mask_qwen3_tools("Qwen/Qwen3-Coder-30B-A3B-Instruct") diff --git a/tests/utils/test_megatron_bridge_utils.py b/tests/utils/test_megatron_bridge_utils.py deleted file mode 100644 index c649293d9..000000000 --- a/tests/utils/test_megatron_bridge_utils.py +++ /dev/null @@ -1,64 +0,0 @@ -import types - -import pytest - -from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_hf_config_for_megatron_bridge - - -@pytest.mark.unit -def test_patch_hf_config_adds_rope_theta_from_rope_parameters(): - hf_config = types.SimpleNamespace(rope_parameters={"rope_theta": 1000000}) - - patched_config = patch_hf_config_for_megatron_bridge(hf_config) - - assert patched_config is hf_config - assert hf_config.rope_theta == 1000000 - - -@pytest.mark.unit -def test_patch_hf_config_does_not_override_existing_rope_theta(): - hf_config = types.SimpleNamespace(rope_theta=500000, rope_parameters={"rope_theta": 1000000}) - - patch_hf_config_for_megatron_bridge(hf_config) - - assert hf_config.rope_theta == 500000 - - -@pytest.mark.unit -def test_patch_hf_config_handles_nested_text_config(): - text_config = types.SimpleNamespace(rope_parameters={"rope_theta": 10000}) - hf_config = types.SimpleNamespace(text_config=text_config) - - patch_hf_config_for_megatron_bridge(hf_config) - - assert text_config.rope_theta == 10000 - - -@pytest.mark.unit -def test_patch_hf_config_handles_pretrained_wrapper_config(): - wrapped_config = types.SimpleNamespace(rope_parameters={"rope_theta": 10000}) - hf_pretrained = types.SimpleNamespace(config=wrapped_config) - - patch_hf_config_for_megatron_bridge(hf_pretrained) - - assert wrapped_config.rope_theta == 10000 - - -@pytest.mark.unit -def test_patch_hf_config_uses_rope_scaling_fallback(): - hf_config = types.SimpleNamespace(rope_scaling={"rope_theta": 10000}) - - patch_hf_config_for_megatron_bridge(hf_config) - - assert hf_config.rope_theta == 10000 - - -@pytest.mark.unit -def test_patch_auto_bridge_hf_config_patches_hf_pretrained(): - hf_config = types.SimpleNamespace(rope_parameters={"rope_theta": 12345}) - bridge = types.SimpleNamespace(hf_pretrained=hf_config) - - patched_bridge = patch_auto_bridge_hf_config(bridge) - - assert patched_bridge is bridge - assert bridge.hf_pretrained.rope_theta == 12345 diff --git a/tests/utils/test_megatron_role_config.py b/tests/utils/test_megatron_role_config.py index 8d6976add..f6ebc5f4b 100644 --- a/tests/utils/test_megatron_role_config.py +++ b/tests/utils/test_megatron_role_config.py @@ -130,28 +130,35 @@ def test_create_training_models_applies_actor_override_without_critic(self, monk args = _base_args(megatron_config_path=path, use_critic=False) class DummyModel: - def __init__(self, model_args): + def __init__(self, model_args, with_ref=False, with_opd_teacher=False): self.args = model_args - self.init_calls = [] + self.with_ref = with_ref + self.with_opd_teacher = with_opd_teacher + self.create_calls = [] self.rollout_manager = None - def async_init(self, model_args, role, with_ref=False, with_opd_teacher=False): - self.args = model_args - self.init_calls.append( + def create(self, rollout_manager=None): + self.rollout_manager = rollout_manager + self.create_calls.append( { - "args": model_args, - "role": role, - "with_ref": with_ref, - "with_opd_teacher": with_opd_teacher, + "args": self.args, + "with_ref": self.with_ref, + "with_opd_teacher": self.with_opd_teacher, + "rollout_manager": rollout_manager, } ) return [7] - def set_rollout_manager(self, rollout_manager): - self.rollout_manager = rollout_manager - - def fake_allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="actor"): - return DummyModel(args) + def fake_allocate_train_group( + args, + num_nodes, + num_gpus_per_node, + pg, + role="actor", + with_ref=False, + with_opd_teacher=False, + ): + return DummyModel(args, with_ref=with_ref, with_opd_teacher=with_opd_teacher) monkeypatch.setattr(placement_group_module, "allocate_train_group", fake_allocate_train_group) monkeypatch.setattr(placement_group_module.ray, "get", lambda value: value) @@ -164,6 +171,9 @@ def fake_allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="acto assert critic_model is None assert actor_model.args.lr == 1e-6 - assert actor_model.init_calls[0]["args"].lr == 1e-6 - assert actor_model.init_calls[0]["role"] == "actor" + assert actor_model.create_calls[0]["args"].lr == 1e-6 assert args.start_rollout_id == 7 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_megatron_server_arguments.py b/tests/utils/test_megatron_server_arguments.py new file mode 100644 index 000000000..d41eb52de --- /dev/null +++ b/tests/utils/test_megatron_server_arguments.py @@ -0,0 +1,88 @@ +import argparse +from argparse import Namespace + +import pytest + +from vime.backends.megatron_utils.server.arguments import ( + add_megatron_server_arguments, + configure_megatron_server_args, + validate_megatron_server_args, +) + +NUM_GPUS = 0 + + +def _server_args(**overrides): + values = dict( + teacher_port=7999, + teacher_warmup_port=7999, + teacher_warmup_timeout_s=3000, + teacher_sample_reduction_chunk_size=4096, + teacher_label_reduction_chunk_size=4096, + megatron_server_max_length=0, + megatron_server_update_timeout_s=3600.0, + megatron_server_warmup=True, + debug_train_only=False, + use_kl_loss=True, + offload_train=True, + use_dynamic_batch_size=True, + use_wandb=True, + kl_coef=0.1, + use_opd=True, + use_critic=True, + keep_old_actor=True, + no_load_optim=False, + no_load_rng=False, + only_train_params_name_list=["actor"], + ) + values.update(overrides) + return Namespace(**values) + + +def test_add_megatron_server_arguments(): + parser = argparse.ArgumentParser() + add_megatron_server_arguments(parser) + + args = parser.parse_args( + [ + "--teacher-port", + "8123", + "--teacher-sample-reduction-chunk-size", + "128", + "--teacher-label-reduction-chunk-size", + "256", + "--no-megatron-server-warmup", + ] + ) + + assert args.teacher_port == 8123 + assert args.teacher_sample_reduction_chunk_size == 128 + assert args.teacher_label_reduction_chunk_size == 256 + assert args.megatron_server_warmup is False + assert args.teacher_warmup_timeout_s == 3000 + assert args.megatron_server_update_timeout_s == 3600.0 + + +def test_configure_megatron_server_args_forces_teacher_only_mode(): + args = configure_megatron_server_args(_server_args()) + validate_megatron_server_args(args) + + assert args.debug_train_only is True + assert args.use_kl_loss is False + assert args.use_opd is False + assert args.use_critic is False + assert args.keep_old_actor is False + assert args.no_load_optim is True + assert args.no_load_rng is True + assert args.only_train_params_name_list == ["nothing_to_train"] + + +def test_validate_megatron_server_args_rejects_invalid_server_values(): + args = configure_megatron_server_args(_server_args(teacher_sample_reduction_chunk_size=0)) + + with pytest.raises(ValueError, match="teacher_sample_reduction_chunk_size"): + validate_megatron_server_args(args) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_npu_accelerator.py b/tests/utils/test_npu_accelerator.py new file mode 100644 index 000000000..709d62955 --- /dev/null +++ b/tests/utils/test_npu_accelerator.py @@ -0,0 +1,170 @@ +"""CPU contracts for NPU selection and pre-Megatron bootstrap ordering.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vime.platforms import current_platform, get_platform, npu, reset_platform_cache +from vime.platforms.npu import NPUAccelerator +from vime.utils import accelerator + + +@pytest.fixture(autouse=True) +def npu_runtime(monkeypatch): + reset_platform_cache() + monkeypatch.setenv("VIME_PLATFORM", "npu") + monkeypatch.delenv("VIME_ACCELERATOR", raising=False) + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("MUSA_PATCH_PATH", raising=False) + monkeypatch.setattr(accelerator, "_REGISTRY", {}) + monkeypatch.setattr(accelerator, "_ACCELERATOR", None) + monkeypatch.setattr(accelerator, "_cuda_available", lambda: False) + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + calls = [] + fake_npu = SimpleNamespace( + is_available=lambda: True, + current_device=lambda: 1, + device_count=lambda: 2, + set_device=lambda index: calls.append(("set_device", index)), + synchronize=lambda: calls.append("synchronize"), + empty_cache=lambda: calls.append("empty_cache"), + ipc_collect=lambda: calls.append("ipc_collect"), + ) + monkeypatch.setattr(torch, "npu", fake_npu, raising=False) + yield calls + reset_platform_cache() + + +def test_platform_registers_npu_before_main_auto_selection(monkeypatch, npu_runtime): + # MindSpeed may also make CUDA's availability probe return true. + monkeypatch.setattr(accelerator, "_cuda_available", lambda: True) + assert current_platform().is_npu + assert isinstance(accelerator.initialize_accelerator(), NPUAccelerator) + assert accelerator.device_type() == "npu" + assert accelerator.process_group_backend() == "hccl" + assert accelerator.process_group_backend("gloo") == "gloo" + assert accelerator.is_accelerator_backend("cpu:gloo,npu:hccl") + assert not accelerator.is_accelerator_backend("gloo") + assert accelerator.distributed_device_id() is None + assert accelerator.visible_devices_env_key() == "ASCEND_RT_VISIBLE_DEVICES" + monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,7") + assert accelerator.resolve_visible_device_id("7") == 1 + accelerator.set_device(1) + accelerator.synchronize() + accelerator.ipc_collect() + accelerator.empty_cache() + assert npu_runtime == [("set_device", 1), "synchronize", "ipc_collect", "empty_cache"] + + +def test_npu_allocator_remains_owned_by_existing_runtime_hooks(monkeypatch): + monkeypatch.setenv("VIME_ENABLE_EXPANDABLE_SEGMENTS", "1") + assert NPUAccelerator().set_allocator_expandable_segments() is False + + +def test_main_npu_override_selects_matching_platform(monkeypatch): + monkeypatch.delenv("VIME_PLATFORM") + monkeypatch.setenv("VIME_ACCELERATOR", "npu") + assert current_platform().is_npu + assert accelerator.get_accelerator().name == "npu" + + +@pytest.mark.parametrize("platform,backend", [("npu", "cuda"), ("cuda", "npu"), ("npu", "musa")]) +def test_conflicting_overrides_fail_before_bootstrap(monkeypatch, platform, backend): + monkeypatch.setenv("VIME_PLATFORM", platform) + monkeypatch.setenv("VIME_ACCELERATOR", backend) + with pytest.raises(ValueError, match="Conflicting VIME_PLATFORM"): + current_platform() + + +def test_registered_npu_does_not_override_explicit_cuda_platform(monkeypatch): + get_platform("npu") + monkeypatch.setenv("VIME_PLATFORM", "cuda") + monkeypatch.setattr(accelerator, "_cuda_available", lambda: True) + monkeypatch.setattr(accelerator.CUDAAccelerator, "is_available", lambda self: True) + assert current_platform().name == "cuda" + assert accelerator.initialize_accelerator().name == "cuda" + + +def test_bootstrap_selects_npu_before_adaptor_and_attention(monkeypatch): + events = [] + bootstrap = current_platform().megatron + monkeypatch.setattr(npu, "_ensure_torch_npu", lambda: events.append("torch_npu")) + monkeypatch.setattr(npu, "_install_safe_empty_cache", lambda: events.append("empty_cache_guard")) + original_import = npu.importlib.import_module + + def import_module(name, *args, **kwargs): + if name in {"megatron_adaptor", "vime.backends.megatron_utils.npu_attention_patch"}: + assert accelerator.get_accelerator().name == "npu" + events.append(name) + bootstrap.bootstrap() # Recursive imports must not repeat initialization. + return SimpleNamespace() + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(npu.importlib, "import_module", import_module) + bootstrap.bootstrap() + bootstrap.bootstrap() + assert events == [ + "torch_npu", + "empty_cache_guard", + "megatron_adaptor", + "vime.backends.megatron_utils.npu_attention_patch", + ] + + +def test_bootstrap_rejects_preselected_cuda_without_replacing_it(monkeypatch): + selected = accelerator.CUDAAccelerator() + monkeypatch.setattr(accelerator, "_ACCELERATOR", selected) + monkeypatch.setattr(npu, "_ensure_torch_npu", lambda: None) + bootstrap = current_platform().megatron + with pytest.raises(RuntimeError, match="already selected 'cuda'"): + bootstrap.bootstrap() + assert accelerator._ACCELERATOR is selected + assert not bootstrap._bootstrapping + assert not bootstrap._bootstrapped + + +def test_repatch_passes_typed_args_and_restores_attention(monkeypatch): + events = [] + full_args = SimpleNamespace(adaptor_default=True) + typed_config = {"weight_nz_mode": 0} + args = SimpleNamespace(vllm_additional_config=typed_config, tensor_model_parallel_size=4) + original_forward = object() + vime_forward = object() + attention_class = SimpleNamespace(forward=vime_forward) + + def apply_features(config): + assert config is full_args + assert config.vllm_additional_config is typed_config + assert config.tensor_model_parallel_size == 4 + assert config.adaptor_default + attention_class.forward = original_forward + events.append("features") + + modules = { + "megatron_adaptor.features_manager.features_manager": SimpleNamespace( + FeaturesManager=SimpleNamespace( + remove_patches=lambda: events.append("remove"), + apply_features_pre_patches=lambda config: events.append(("pre", config)), + apply_features_patches=apply_features, + ) + ), + "megatron_adaptor.utils.args_utils": SimpleNamespace(get_full_args=lambda: full_args), + "vime.backends.megatron_utils.npu_attention_patch": SimpleNamespace( + DotProductAttention=attention_class, + npu_dot_product_attention_forward=vime_forward, + ), + } + bootstrap = current_platform().megatron + original_import = npu.importlib.import_module + + def import_module(name, *args, **kwargs): + if name in modules: + return modules[name] + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(npu.importlib, "import_module", import_module) + bootstrap.repatch(args) + assert events == ["remove", ("pre", full_args), "features"] + assert attention_class.forward is vime_forward + assert args.vllm_additional_config is typed_config diff --git a/tests/utils/test_platform_contract.py b/tests/utils/test_platform_contract.py new file mode 100644 index 000000000..c930fed82 --- /dev/null +++ b/tests/utils/test_platform_contract.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import sys +from argparse import Namespace +from types import ModuleType, SimpleNamespace + +import pytest + +from vime.platforms import current_platform, reset_platform_cache + + +@pytest.fixture(autouse=True) +def _reset_platform_selection(): + reset_platform_cache() + yield + reset_platform_cache() + + +def test_vime_platform_override_selects_cuda(monkeypatch): + monkeypatch.setenv("VIME_PLATFORM", "cuda") + + platform = current_platform() + + assert platform.name == "cuda" + assert platform.ray.resource_name == "GPU" + assert platform.ray.visible_devices_env == "CUDA_VISIBLE_DEVICES" + assert platform.checkpoint.default_megatron_to_hf_mode == "raw" + + +def test_vime_platform_override_selects_npu_without_vendor_import(monkeypatch): + monkeypatch.setenv("VIME_PLATFORM", "npu") + before = {name for name in ("torch_npu", "vllm_ascend", "mindspeed") if name in sys.modules} + + platform = current_platform() + + after = {name for name in ("torch_npu", "vllm_ascend", "mindspeed") if name in sys.modules} + assert platform.name == "npu" + assert platform.ray.resource_name == "NPU" + assert platform.checkpoint.default_megatron_to_hf_mode == "bridge" + assert before == after + + +def test_unknown_explicit_platform_has_clear_error(monkeypatch): + monkeypatch.setenv("VIME_PLATFORM", "not-registered") + + with pytest.raises(ValueError, match="Unknown Vime platform"): + current_platform() + + +def test_npu_auto_detection_does_not_import_vendor_without_device_nodes(monkeypatch): + from vime.platforms import npu + + monkeypatch.setattr(npu.os.path, "exists", lambda path: False) + monkeypatch.setattr(npu, "glob", lambda pattern: []) + + def fail_import(name): + raise AssertionError(f"unexpected import during negative NPU detection: {name}") + + monkeypatch.setattr(npu.importlib, "import_module", fail_import) + + assert npu.detect_npu() is False + + +def test_npu_safe_empty_cache_wraps_original_once(monkeypatch): + from vime.platforms import npu + + calls = [] + + def original_empty_cache(): + calls.append("original") + raise RuntimeError("allocator is between offload states") + + fake_torch = SimpleNamespace( + npu=SimpleNamespace(empty_cache=original_empty_cache), + cuda=SimpleNamespace(empty_cache=lambda: None), + ) + monkeypatch.setattr(npu.importlib, "import_module", lambda name: fake_torch if name == "torch" else None) + + npu._install_safe_empty_cache() + wrapped = fake_torch.npu.empty_cache + wrapped() + npu._install_safe_empty_cache() + + assert calls == ["original"] + assert fake_torch.npu.empty_cache is wrapped + assert fake_torch.cuda.empty_cache is wrapped + + +@pytest.mark.parametrize( + ("name", "bundle", "actor_options"), + [ + ("cuda", {"GPU": 2, "CPU": 3}, {"num_gpus": 0.4}), + ("npu", {"NPU": 2, "CPU": 3}, {"resources": {"NPU": 0.4}}), + ], +) +def test_ray_resource_contract(monkeypatch, name, bundle, actor_options): + monkeypatch.setenv("VIME_PLATFORM", name) + ray_spec = current_platform().ray + + assert ray_spec.bundle_resources(device_count=2, cpu_count=3) == bundle + assert ray_spec.actor_options(0.4) == actor_options + + +def test_npu_runtime_env_is_scoped_to_npu_provider(monkeypatch, tmp_path): + toolkit = tmp_path / "toolkit" + (toolkit / "python" / "site-packages" / "acl").mkdir(parents=True) + monkeypatch.setenv("ASCEND_TOOLKIT_HOME", str(toolkit)) + monkeypatch.setenv("VIME_PLATFORM", "npu") + args = Namespace(offload_train=True, train_backend="megatron", colocate=True) + + train_env = current_platform().ray.train_runtime_env(args, {"BASE": "1"}) + rollout_env = current_platform().ray.rollout_runtime_env(args, {"BASE": "1"}) + + assert train_env["TMS_HOOK_MODE"] == "torch" + assert train_env["TMS_REGION_TAG"] == "training" + assert train_env["TMS_ENABLE_CPU_BACKUP"] == "1" + assert train_env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False" + assert str(toolkit / "python" / "site-packages") in train_env["PYTHONPATH"] + assert rollout_env["VLLM_USE_AOT_COMPILE"] == "0" + assert rollout_env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False" + + +def test_npu_vllm_env_replaces_cuda_and_rocm_visibility(monkeypatch): + monkeypatch.setenv("VIME_PLATFORM", "npu") + platform = current_platform() + + env = platform.vllm.subprocess_env( + { + "KEEP": "1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "CUDA_VISIBLE_DEVICES": "0,1", + "HIP_VISIBLE_DEVICES": "0,1", + }, + visible_devices="4,5", + colocate=True, + ) + + assert env["KEEP"] == "1" + assert "PYTORCH_CUDA_ALLOC_CONF" not in env + assert "CUDA_VISIBLE_DEVICES" not in env + assert "HIP_VISIBLE_DEVICES" not in env + assert env["ASCEND_RT_VISIBLE_DEVICES"] == "4,5" + assert env["PYTORCH_NPU_ALLOC_CONF"] == "expandable_segments:False" + + +@pytest.mark.parametrize("name,backends", [("cuda", ("nccl", "ipc")), ("npu", ("hccl", "npu_ipc"))]) +def test_weight_transfer_provider_selects_matching_backend_and_init_info(monkeypatch, name, backends): + from vime.platforms import get_platform + + plugin_calls = [] + for colocate, backend in zip((False, True), backends, strict=True): + if name == "cuda": + module_name = f"vllm.distributed.weight_transfer.{backend}_engine" + cls_name = "IPCTrainerInitInfo" if colocate else "NCCLTrainerInitInfo" + else: + module_name = f"vllm_ascend.distributed.weight_transfer.{backend}_engine" + cls_name = "NPUIPCTrainerInitInfo" if colocate else "HCCLTrainerInitInfo" + module = ModuleType(module_name) + info_cls = type(cls_name, (SimpleNamespace,), {"backend": backend}) + setattr(module, cls_name, info_cls) + monkeypatch.setitem(sys.modules, module_name, module) + + plugins = ModuleType("vllm.plugins") + plugins.load_general_plugins = lambda: plugin_calls.append("load") + monkeypatch.setitem(sys.modules, "vllm.plugins", plugins) + monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) + + ops = get_platform(name).weight_transfer + info = ops.trainer_init_info(colocate=colocate, rank=3, packed=True) + assert ops.backend("ipc" if colocate else "nccl") == info.backend == backend + assert isinstance(info, info_cls) + assert (info.rank, info.packed) == (3, True) + + assert plugin_calls == (["load", "load"] if name == "npu" else []) + for backend in ("npu_ipc", "hccl", "custom"): + assert ops.backend(backend) == backend + + +@pytest.mark.parametrize("platform", ["cuda", "npu"]) +@pytest.mark.parametrize("chained", [False, True]) +def test_optimizer_state_initialization_reuses_megatron_callback(monkeypatch, platform, chained): + monkeypatch.setenv("VIME_PLATFORM", platform) + calls = [] + + def init_state(optimizer, config): + calls.append((optimizer, config)) + + optimizers = [ + SimpleNamespace(optimizer=object(), config=object(), init_state_fn=init_state) + for _ in range(2 if chained else 1) + ] + optimizer = SimpleNamespace(chained_optimizers=optimizers) if chained else optimizers[0] + + current_platform().megatron.initialize_optimizer_state(optimizer) + + expected = [(opt.optimizer, opt.config) for opt in optimizers] if platform == "npu" else [] + assert calls == expected + + +@pytest.mark.parametrize("empty_optimizer", [False, True]) +def test_npu_optimizer_state_initialization_skips_missing_state_or_optimizer(monkeypatch, empty_optimizer): + monkeypatch.setenv("VIME_PLATFORM", "npu") + + def unexpected_init(*args): + raise AssertionError("an empty optimizer must not initialize state") + + optimizer = SimpleNamespace( + optimizer=None if empty_optimizer else object(), + config=object(), + init_state_fn=unexpected_init if empty_optimizer else None, + ) + + current_platform().megatron.initialize_optimizer_state(optimizer) + + +@pytest.mark.parametrize("platform", ["cuda", "npu"]) +def test_eval_only_has_no_optimizer_state_to_initialize(monkeypatch, platform): + monkeypatch.setenv("VIME_PLATFORM", platform) + current_platform().megatron.initialize_optimizer_state(None) + + +def test_memory_utils_keep_main_accelerator_surface(monkeypatch): + from vime.utils import memory_utils + + calls = [] + fake_accelerator = SimpleNamespace( + synchronize=lambda: calls.append("synchronize"), + empty_cache=lambda: calls.append("empty_cache"), + ) + fake_torch = SimpleNamespace(_C=SimpleNamespace(_host_emptyCache=lambda: calls.append("empty_host_cache"))) + monkeypatch.setattr(memory_utils, "accelerator", fake_accelerator) + monkeypatch.setattr(memory_utils, "torch", fake_torch) + monkeypatch.setattr(memory_utils.gc, "collect", lambda: calls.append("gc")) + + memory_utils.clear_memory(clear_host_memory=True) + + assert calls == ["synchronize", "gc", "empty_cache", "empty_host_cache"] diff --git a/tests/utils/test_ray_platform_integration.py b/tests/utils/test_ray_platform_integration.py new file mode 100644 index 000000000..e5d19dcf4 --- /dev/null +++ b/tests/utils/test_ray_platform_integration.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + + +def _fake_platform(ray_ops, *, is_npu=False): + return SimpleNamespace(ray=ray_ops, is_npu=is_npu) + + +@pytest.mark.parametrize( + ("platform_name", "visible_env", "assigned_id", "expected_resource"), + [ + ("cuda", "CUDA_VISIBLE_DEVICES", "7", "GPU"), + ("npu", "ASCEND_RT_VISIBLE_DEVICES", "9", "NPU"), + ], +) +def test_platform_ray_accelerator_ids_and_local_mapping( + monkeypatch, + platform_name, + visible_env, + assigned_id, + expected_resource, +): + import ray + + from vime.platforms import get_platform + + platform = get_platform(platform_name) + monkeypatch.setenv(visible_env, f"unused,{assigned_id}") + if platform_name == "cuda": + monkeypatch.setattr(ray, "get_gpu_ids", lambda: [assigned_id]) + else: + context = SimpleNamespace(get_accelerator_ids=lambda: {"NPU": [assigned_id]}) + monkeypatch.setattr(ray, "get_runtime_context", lambda: context) + + assert platform.ray.resource_name == expected_resource + assert platform.ray.accelerator_ids() == [assigned_id] + assert platform.ray.local_device_id() == 1 + + +def test_placement_group_uses_platform_ray_resource_contract(monkeypatch): + from vime.ray import placement_group as placement_group_module + + ray_ops = SimpleNamespace( + resource_name="ACCEL", + bundle_resources=Mock(side_effect=lambda: {"ACCEL": 1, "CPU": 1}), + actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}), + ) + monkeypatch.setattr(placement_group_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=True)) + + created = {} + + class FakePlacementGroup: + def ready(self): + return "ready" + + def fake_placement_group(bundles, strategy): + created["bundles"] = bundles + created["strategy"] = strategy + return FakePlacementGroup() + + actor_options = [] + + class FakeInfoActor: + def __init__(self, result): + self.get_ip_and_gpu_id = SimpleNamespace(remote=lambda: result) + + class FakeInfoActorClass: + results = iter([("10.0.0.1", "3"), ("10.0.0.1", "1")]) + + @classmethod + def options(cls, **options): + actor_options.append(options) + result = next(cls.results) + return SimpleNamespace(remote=lambda: FakeInfoActor(result)) + + monkeypatch.setattr(placement_group_module, "placement_group", fake_placement_group) + monkeypatch.setattr(placement_group_module, "InfoActor", FakeInfoActorClass) + wait_results = iter([([], ["ready"]), (["ready"], [])]) + monkeypatch.setattr(placement_group_module.ray, "wait", lambda *_args, **_kwargs: next(wait_results)) + monkeypatch.setattr(placement_group_module.ray, "cluster_resources", lambda: {"ACCEL": 2}) + monkeypatch.setattr(placement_group_module.ray, "available_resources", lambda: {"ACCEL": 2}) + monkeypatch.setattr(placement_group_module.ray, "get", lambda value: value) + monkeypatch.setattr(placement_group_module.ray, "kill", lambda _actor: None) + + pg, reordered_indices, reordered_ids = placement_group_module._create_placement_group(2) + + assert isinstance(pg, FakePlacementGroup) + assert created == { + "bundles": [{"ACCEL": 1, "CPU": 1}, {"ACCEL": 1, "CPU": 1}], + "strategy": "PACK", + } + assert reordered_indices == [1, 0] + assert reordered_ids == ["1", "3"] + assert [options["resources"] for options in actor_options] == [{"ACCEL": 1}, {"ACCEL": 1}] + assert [options["num_gpus"] for options in actor_options] == [0, 0] + assert ray_ops.bundle_resources.call_count == 2 + assert [entry.args for entry in ray_ops.actor_options.call_args_list] == [(1,), (1,)] + + +def test_ray_noset_visible_devices_keeps_ascend_entry(): + from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST + + assert "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES" in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST + + +def test_train_group_uses_npu_runtime_env_and_actor_resources(monkeypatch): + from vime.ray import actor_group as actor_group_module + + runtime_env_inputs = [] + ray_ops = SimpleNamespace( + train_runtime_env=lambda args, env: runtime_env_inputs.append((args, dict(env))) + or {**env, "PLATFORM_ENV": "1"}, + actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}), + ) + monkeypatch.setattr(actor_group_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=True)) + + actor_module = types.ModuleType("vime.backends.megatron_utils.actor") + actor_module.MegatronTrainRayActor = object + monkeypatch.setitem(sys.modules, "vime.backends.megatron_utils.actor", actor_module) + + remote_declarations = [] + actor_allocations = [] + + class FakeActorHandle: + get_master_addr_and_port = SimpleNamespace(remote=lambda: ("127.0.0.1", 20000)) + init = SimpleNamespace(remote=lambda *_args, **_kwargs: 0) + + class FakeRemoteActor: + def options(self, **options): + actor_allocations.append(options) + return self + + def remote(self, *_args): + return FakeActorHandle() + + def fake_remote(**options): + remote_declarations.append(options) + return lambda _actor_impl: FakeRemoteActor() + + monkeypatch.setattr(actor_group_module.ray, "remote", fake_remote) + monkeypatch.setattr(actor_group_module.ray, "get", lambda value: value) + + args = SimpleNamespace( + train_env_vars={"USER_ENV": "yes"}, + offload_train=True, + train_backend="megatron", + colocate=False, + use_routing_replay=False, + ) + group = actor_group_module.RayTrainGroup( + args=args, + num_nodes=1, + num_gpus_per_node=1, + pg=(object(), [0], [7]), + num_gpus_per_actor=0.4, + ) + assert group.create() == [0] + + assert runtime_env_inputs[0][0] is args + assert runtime_env_inputs[0][1]["USER_ENV"] == "yes" + assert remote_declarations[0]["runtime_env"]["env_vars"]["PLATFORM_ENV"] == "1" + assert remote_declarations[0]["num_gpus"] == 1 + assert actor_allocations[0]["resources"] == {"ACCEL": 0.4} + assert actor_allocations[0]["num_gpus"] == 0 + ray_ops.actor_options.assert_called_once_with(0.4) + + +@pytest.mark.parametrize("is_npu", [False, True]) +def test_rollout_engine_delegates_runtime_env_and_actor_options(monkeypatch, is_npu): + from vime.backends.vllm_utils import engine_group as rollout_module + + runtime_env_inputs = [] + ray_ops = SimpleNamespace( + rollout_runtime_env=lambda args, env: runtime_env_inputs.append((args, dict(env))) + or {**env, "PLATFORM_ENV": "rollout"}, + actor_options=Mock(side_effect=lambda fraction: {"resources": {"ACCEL": fraction}}), + ) + monkeypatch.setattr(rollout_module, "current_platform", lambda: _fake_platform(ray_ops, is_npu=is_npu)) + + actor_options = [] + + class FakeEngine: + init = SimpleNamespace(remote=lambda **_kwargs: "init-ref") + + class FakeRemoteActor: + def options(self, **options): + actor_options.append(options) + return self + + def remote(self, *_args, **_kwargs): + return FakeEngine() + + monkeypatch.setattr(rollout_module.ray, "remote", lambda _actor_impl: FakeRemoteActor()) + monkeypatch.setattr( + rollout_module, + "_allocate_rollout_engine_addr_and_ports_normal", + lambda **_kwargs: ({0: {}}, {0: 15001}), + ) + + args = SimpleNamespace( + debug_train_only=False, + num_gpus_per_node=8, + rollout_num_gpus=1, + rollout_num_gpus_per_engine=1, + rollout_external=False, + colocate=True, + ) + group = rollout_module.ServerGroup( + args=args, + pg=(object(), [0], [7]), + all_engines=[None], + num_gpus_per_engine=1, + num_new_engines=1, + ) + + handles, cursors = group.start_engines() + + assert handles == ["init-ref"] + assert cursors == {0: 15001} + if is_npu: + assert runtime_env_inputs[0][0] is args + assert actor_options[0]["runtime_env"]["env_vars"]["PLATFORM_ENV"] == "rollout" + assert actor_options[0]["resources"] == {"ACCEL": 0.2} + assert actor_options[0]["num_gpus"] == 0 + ray_ops.actor_options.assert_called_once_with(0.2) + else: + assert runtime_env_inputs == [] + assert "resources" not in actor_options[0] + assert "PLATFORM_ENV" not in actor_options[0]["runtime_env"]["env_vars"] + assert actor_options[0]["num_gpus"] == 0.2 + ray_ops.actor_options.assert_not_called() + + # The main placement check must still run before an invalid slot is used. + group.all_engines = [None] + group.gpu_offset = 1 + with pytest.raises(ValueError, match="Invalid rollout server group GPU placement"): + group.start_engines() + + +def test_train_actor_uses_npu_local_device_mapping(monkeypatch): + from vime.ray import train_actor as train_actor_module + + local_device_id = Mock(return_value=5) + monkeypatch.setattr( + train_actor_module, + "current_platform", + lambda: _fake_platform(SimpleNamespace(local_device_id=local_device_id), is_npu=True), + ) + + assert train_actor_module.get_local_gpu_id() == 5 + local_device_id.assert_called_once_with() + + +def test_train_actor_keeps_main_cuda_local_device_mapping(monkeypatch): + from vime.ray import train_actor as train_actor_module + + monkeypatch.setattr( + train_actor_module, + "current_platform", + lambda: _fake_platform(SimpleNamespace(), is_npu=False), + ) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,3") + monkeypatch.setattr( + train_actor_module.accelerator, "_ACCELERATOR", train_actor_module.accelerator.CUDAAccelerator() + ) + monkeypatch.setattr(train_actor_module.ray, "get_gpu_ids", lambda: ["3"]) + + assert train_actor_module.get_local_gpu_id() == 1 diff --git a/tests/utils/test_update_weight_from_distributed.py b/tests/utils/test_update_weight_from_distributed.py new file mode 100644 index 000000000..f63a72b2b --- /dev/null +++ b/tests/utils/test_update_weight_from_distributed.py @@ -0,0 +1,462 @@ +"""CPU unit tests for the vLLM trainer-side weight-transfer adapter.""" + +from __future__ import annotations + +import importlib +import inspect +import sys +import types +from dataclasses import dataclass, field +from pathlib import Path + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs +import pytest +import torch + +from vime.utils.types import ParamInfo + +MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_distributed" +COMMON_MODULE = "vime.backends.megatron_utils.update_weight.common" +DIRECT_MODULE = "vime.backends.megatron_utils.update_weight.hf_weight_iterator_direct" +CONVERTER_MODULE = "vime.backends.megatron_utils.megatron_to_hf" + + +@pytest.fixture(scope="module") +def update_module(): + module_names = ( + "megatron", + "megatron.core", + "megatron.core.parallel_state", + "megatron.core.transformer", + "megatron.core.transformer.transformer_layer", + "ray", + "ray.actor", + "vime.utils.distributed_utils", + COMMON_MODULE, + MODULE_PATH, + ) + saved = _unit_stubs.save_sys_modules(module_names) + for name in module_names: + sys.modules.pop(name, None) + _unit_stubs.install_megatron_mpu_stub() + _unit_stubs.install_ray_stub() + _unit_stubs.install_vime_distributed_utils_stub() + try: + yield importlib.import_module(MODULE_PATH) + finally: + _unit_stubs.restore_sys_modules(saved) + + +@dataclass +class RemoteCall: + args: tuple + kwargs: dict + + +class RecordingRemoteMethod: + def __init__(self): + self.calls: list[RemoteCall] = [] + + def remote(self, *args, **kwargs): + self.calls.append(RemoteCall(args, kwargs)) + return "ref" + + +@dataclass +class RecordingEngine: + init_weight_transfer_engine: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + start_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + start_draft_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + update_weights: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + finish_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + pause_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + flush_cache: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + continue_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + post_process_weights: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + + +class RecordingTrainer: + def __init__(self, client, *, source=None, fail=False): + self.client = client + self.source = source + self.fail = fail + self.draft_states = [] + self.source_draft_states = [] + self.shutdown_calls = 0 + + def send_weights(self): + self.draft_states.append(self.client.draft) + self.source_draft_states.append(getattr(self.source, "draft", False)) + self.client.start_weight_update() + if self.fail: + raise RuntimeError("transfer failed") + self.client.update_weights({"names": []}) + self.client.finish_weight_update() + + def shutdown(self): + self.shutdown_calls += 1 + + +@pytest.mark.unit +def test_ray_client_fans_out_and_offsets_nccl_ranks(update_module): + engines = [RecordingEngine(), RecordingEngine()] + client = update_module.VimeRayWeightSyncClient(engines, lambda: 7, [2, 4]) + + client.init_weight_transfer_engine({"rank_offset": 1, "world_size": 7}) + client.start_weight_update() + client.update_weights({"names": ["weight"]}) + client.finish_weight_update() + + assert engines[0].init_weight_transfer_engine.calls[0].args[0]["init_info"]["rank_offset"] == 1 + assert engines[1].init_weight_transfer_engine.calls[0].args[0]["init_info"]["rank_offset"] == 3 + assert len(engines[0].start_weight_update.calls) == 1 + assert len(engines[1].update_weights.calls) == 1 + assert engines[0].finish_weight_update.calls[0].kwargs == {"weight_version": "7"} + + +@pytest.mark.unit +def test_ray_client_selects_draft_lifecycle(update_module): + engine = RecordingEngine() + client = update_module.VimeRayWeightSyncClient([engine], lambda: 1) + client.draft = True + + client.start_weight_update() + + assert engine.start_weight_update.calls == [] + assert len(engine.start_draft_weight_update.calls) == 1 + + +@pytest.mark.unit +def test_weight_source_caches_metadata_and_reiterates(update_module, monkeypatch): + class ParamMeta: + def __init__(self, name, dtype, shape): + self.name = name + self.dtype = dtype + self.shape = shape + + base_module = types.ModuleType("vllm.distributed.weight_transfer.base") + base_module.ParamMeta = ParamMeta + monkeypatch.setitem(sys.modules, "vllm.distributed.weight_transfer.base", base_module) + + calls = [] + + class Iterator: + def get_hf_weight_chunks(self, weights): + calls.append(weights) + yield [("a", torch.zeros(2)), ("b", torch.ones(3))] + + source = update_module.HfWeightSource(Iterator(), lambda: {"version": len(calls)}) + + assert [item.name for item in source.metadata()] == ["a", "b"] + assert [item.name for item in source.metadata()] == ["a", "b"] + assert [name for name, _ in source] == ["a", "b"] + assert len(calls) == 2 + + +@pytest.mark.unit +def test_weight_source_switches_to_draft_weights(update_module, monkeypatch): + class ParamMeta: + def __init__(self, name, dtype, shape): + self.name = name + self.dtype = dtype + self.shape = shape + + base_module = types.ModuleType("vllm.distributed.weight_transfer.base") + base_module.ParamMeta = ParamMeta + monkeypatch.setitem(sys.modules, "vllm.distributed.weight_transfer.base", base_module) + + class Iterator: + def get_hf_weight_chunks(self, weights): + yield [("policy", weights["policy"])] + + source = update_module.HfWeightSource( + Iterator(), + lambda: {"policy": torch.zeros(1)}, + lambda: [("draft", torch.ones(2))], + ) + + assert [name for name, _ in source] == ["policy"] + source.draft = True + assert [item.name for item in source.metadata()] == ["draft"] + assert [name for name, _ in source] == ["draft"] + + +@pytest.mark.unit +def test_nccl_trainer_uses_single_packed_buffer(update_module, monkeypatch): + from vime.platforms import get_platform + + adapter = sys.modules[update_module.create_nccl_trainer.__module__] + monkeypatch.setattr(adapter, "current_platform", lambda: get_platform("cuda")) + created = [] + + class NCCLTrainerInitInfo: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class Factory: + @staticmethod + def trainer_init(init_info, *, client, source): + created.append((init_info, client, source)) + return "trainer" + + factory_module = types.ModuleType("vllm.distributed.weight_transfer.factory") + factory_module.WeightTransferTrainerFactory = Factory + nccl_module = types.ModuleType("vllm.distributed.weight_transfer.nccl_engine") + nccl_module.NCCLTrainerInitInfo = NCCLTrainerInitInfo + monkeypatch.setitem(sys.modules, factory_module.__name__, factory_module) + monkeypatch.setitem(sys.modules, nccl_module.__name__, nccl_module) + monkeypatch.setattr(adapter.dist, "get_rank", lambda: 0) + monkeypatch.setattr(adapter.dist, "broadcast_object_list", lambda *args, **kwargs: None) + monkeypatch.setattr(adapter, "get_gloo_group", lambda: None) + monkeypatch.setattr( + sys.modules["ray"], + "_private", + types.SimpleNamespace(services=types.SimpleNamespace(get_node_ip_address=lambda: "127.0.0.1")), + raising=False, + ) + + assert adapter.create_nccl_trainer("client", "source", [2, 2]) == "trainer" + init_info, client, source = created[0] + assert init_info.packed_num_buffers == 1 + assert init_info.world_size == 5 + assert client == "client" + assert source == "source" + + +@pytest.mark.unit +def test_connect_replaces_existing_trainer(update_module, monkeypatch): + updater = object.__new__(update_module.UpdateWeightFromDistributed) + updater.args = types.SimpleNamespace(rollout_num_gpus_per_engine=2) + updater.weight_version = 0 + updater._source = object() + old_trainer = RecordingTrainer(object()) + updater._trainer = old_trainer + engine = RecordingEngine() + created = [] + + def create_trainer(client, source, gpu_counts): + created.append((client, source, gpu_counts)) + return RecordingTrainer(client) + + monkeypatch.setattr(update_module, "create_nccl_trainer", create_trainer) + updater.connect_rollout_engines([engine], object(), engine_gpu_counts=[4]) + + assert old_trainer.shutdown_calls == 1 + assert created[0][1:] == (updater._source, [4]) + assert updater._trainer is not old_trainer + + +def _updater_for_transfer(update_module, *, mtp=False, dspark=False, fail=False): + updater = object.__new__(update_module.UpdateWeightFromDistributed) + updater.args = types.SimpleNamespace( + enable_mtp_training=mtp, + dspark_enabled=dspark, + vllm_speculative_config={"method": "mtp"} if mtp else {"method": "dspark"} if dspark else None, + ) + updater.quantization_config = None + updater.weight_version = 0 + updater.update_weight_metrics = {} + updater.rollout_engines = [RecordingEngine()] + client = update_module.VimeRayWeightSyncClient(updater.rollout_engines, lambda: updater.weight_version) + updater._source = types.SimpleNamespace(draft=False) + updater._trainer = RecordingTrainer(client, source=updater._source, fail=fail) + return updater + + +@pytest.mark.unit +def test_update_uses_native_main_and_draft_lifecycles(update_module, monkeypatch): + updater = _updater_for_transfer(update_module, mtp=True) + monkeypatch.setattr(update_module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + updater.update_weights() + + engine = updater.rollout_engines[0] + assert updater._trainer.draft_states == [False, True] + assert len(engine.pause_generation.calls) == 1 + assert len(engine.flush_cache.calls) == 1 + assert len(engine.start_weight_update.calls) == 1 + assert len(engine.start_draft_weight_update.calls) == 1 + assert len(engine.finish_weight_update.calls) == 2 + assert len(engine.continue_generation.calls) == 1 + + +@pytest.mark.unit +def test_dspark_update_uses_nccl_draft_lifecycle(update_module, monkeypatch): + updater = _updater_for_transfer(update_module, dspark=True) + monkeypatch.setattr(update_module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + updater.update_weights() + + engine = updater.rollout_engines[0] + assert updater._trainer.draft_states == [False, True] + assert updater._trainer.source_draft_states == [False, True] + assert updater._source.draft is False + assert len(engine.start_draft_weight_update.calls) == 1 + + +@pytest.mark.unit +def test_failed_transfer_does_not_resume_generation(update_module, monkeypatch): + updater = _updater_for_transfer(update_module, fail=True) + monkeypatch.setattr(update_module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + with pytest.raises(RuntimeError, match="transfer failed"): + updater.update_weights() + + assert updater.rollout_engines[0].continue_generation.calls == [] + + +@pytest.fixture +def weight_modules(): + module_names = ( + "megatron", + "megatron.core", + "megatron.core.parallel_state", + "megatron.core.transformer", + "megatron.core.transformer.transformer_layer", + CONVERTER_MODULE, + COMMON_MODULE, + DIRECT_MODULE, + ) + saved = _unit_stubs.save_sys_modules(module_names) + for name in module_names: + sys.modules.pop(name, None) + _unit_stubs.install_megatron_mpu_stub() + converter = types.ModuleType(CONVERTER_MODULE) + converter.convert_to_hf = lambda *args, **kwargs: [] + sys.modules[CONVERTER_MODULE] = converter + try: + yield importlib.import_module(COMMON_MODULE), importlib.import_module(DIRECT_MODULE) + finally: + _unit_stubs.restore_sys_modules(saved) + + +class Handle: + def wait(self): + pass + + +@pytest.mark.unit +def test_vllm_weight_iterator_keeps_checkpoint_scale_layout(weight_modules): + _, direct_module = weight_modules + + assert inspect.signature(direct_module.HfWeightIteratorDirect).parameters["transform_ue8m0"].default is False + + +def _param_info(name: str, param: torch.Tensor, src_rank: int = 0) -> ParamInfo: + return ParamInfo(name, param.dtype, param.shape, {}, param.nbytes, src_rank) + + +def _tp_param(values, partition_dim: int) -> torch.nn.Parameter: + parameter = torch.nn.Parameter(torch.tensor(values, dtype=torch.float32)) + parameter.tensor_model_parallel = True + parameter.partition_dim = partition_dim + parameter.partition_stride = 1 + return parameter + + +@pytest.mark.unit +def test_single_tp_returns_parameter_without_collective(monkeypatch, weight_modules): + common, _ = weight_modules + parameter = _tp_param([[1.0, 1.0]], partition_dim=0) + calls = [] + monkeypatch.setattr(common.mpu, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(common.dist, "all_gather", lambda *args, **kwargs: calls.append(args)) + + gathered = common.all_gather_param("decoder.weight", parameter) + + assert gathered.data_ptr() == parameter.data_ptr() + assert calls == [] + + +@pytest.mark.unit +def test_all_gather_params_async_restores_layouts(monkeypatch, weight_modules): + common, _ = weight_modules + direct = torch.nn.Parameter(torch.tensor([99.0])) + direct.tensor_model_parallel = False + column = _tp_param([[1.0, 2.0], [3.0, 4.0]], partition_dim=0) + glu = _tp_param([[1.0], [2.0], [10.0], [20.0]], partition_dim=0) + glu.partition_stride = 2 + row = _tp_param([[1.0, 2.0], [3.0, 4.0]], partition_dim=0) + entries = [ + (_param_info("dense", direct), direct), + (_param_info("linear.weight", column), column), + (_param_info("linear_fc1.weight", glu), glu), + (_param_info("linear_fc2.weight", row), row), + ] + remote_parts = iter( + [ + torch.tensor([[5.0, 6.0], [7.0, 8.0]]), + torch.tensor([[3.0], [4.0], [30.0], [40.0]]), + torch.tensor([[5.0, 6.0], [7.0, 8.0]]), + ] + ) + calls = [] + + def all_gather(partitions, local, group, async_op): + calls.append((group, async_op)) + partitions[0].copy_(local) + partitions[1].copy_(next(remote_parts)) + return Handle() + + monkeypatch.setattr(common.mpu, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(common.mpu, "get_tensor_model_parallel_group", lambda: "tp") + monkeypatch.setattr(common.dist, "all_gather", all_gather) + + gathered = common.all_gather_params_async(entries) + + assert calls == [("tp", True)] * 3 + assert gathered[0].data_ptr() == direct.data_ptr() + assert torch.equal(gathered[1], torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])) + assert torch.equal(gathered[2], torch.tensor([[1.0], [2.0], [3.0], [4.0], [10.0], [20.0], [30.0], [40.0]])) + assert torch.equal(gathered[3], torch.tensor([[1.0, 2.0, 5.0, 6.0], [3.0, 4.0, 7.0, 8.0]])) + + +@pytest.mark.unit +def test_broadcast_expert_params_coalesces_by_source(monkeypatch, weight_modules): + _, direct = weight_modules + params = [torch.tensor([float(index)]) for index in range(4)] + infos = [ + _param_info("layers.0.experts.0.weight", params[0], src_rank=4), + _param_info("layers.0.experts.1.weight", params[1], src_rank=5), + _param_info("layers.0.dense.weight", params[2], src_rank=4), + _param_info("layers.0.experts.2.weight", params[3], src_rank=9), + ] + calls = [] + monkeypatch.setattr(direct.mpu, "get_expert_model_parallel_group", lambda: "ep") + monkeypatch.setattr( + direct.dist, + "_broadcast_coalesced", + lambda group, tensors, buffer_size, src: calls.append((group, tensors, buffer_size, src)), + ) + + direct._broadcast_expert_params(infos, params, 1024, {4: 0, 5: 1, 9: 0}) + + assert calls == [ + ("ep", [params[0], params[3]], 1024, 0), + ("ep", [params[1]], 1024, 1), + ] + + +@pytest.mark.unit +def test_ep_broadcast_source_map_tracks_pp_groups(monkeypatch, weight_modules): + _, direct = weight_modules + monkeypatch.setattr(direct.mpu, "get_expert_model_parallel_group", lambda: "ep") + monkeypatch.setattr(direct.mpu, "get_expert_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(direct.mpu, "get_pipeline_model_parallel_group", lambda: "pp") + monkeypatch.setattr(direct.dist, "get_process_group_ranks", lambda group: [0, 2] if group == "pp" else [0, 1]) + + def all_gather_object(output, local_pp_group, group): + assert local_pp_group == [0, 2] + assert group == "ep" + output[:] = [[0, 2], [1, 3]] + + monkeypatch.setattr(direct.dist, "all_gather_object", all_gather_object) + + assert direct._get_ep_broadcast_src_rank_map() == {0: 0, 2: 0, 1: 1, 3: 1} diff --git a/tests/utils/test_update_weight_from_tensor.py b/tests/utils/test_update_weight_from_tensor.py new file mode 100644 index 000000000..0e8ac0d61 --- /dev/null +++ b/tests/utils/test_update_weight_from_tensor.py @@ -0,0 +1,398 @@ +"""CPU unit tests for native and rank-local colocated weight transfer.""" + +from __future__ import annotations + +import importlib +import sys +import types +from argparse import Namespace +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import MagicMock + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs +import pytest +import torch + +from vime.utils.types import ParamInfo + +MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_tensor" +COMMON_MODULE = "vime.backends.megatron_utils.update_weight.common" +HF_DIRECT_MODULE = "vime.backends.megatron_utils.update_weight.hf_weight_iterator_direct" +DISTRIBUTED_MODULE = "vime.backends.megatron_utils.update_weight.update_weight_from_distributed" + + +@pytest.fixture(scope="module") +def update_module(): + import torch.distributed as torch_dist + + module_names = ( + "megatron", + "megatron.core", + "megatron.core.parallel_state", + "megatron.core.transformer", + "megatron.core.transformer.transformer_layer", + "ray", + "ray.actor", + "vime.utils.distributed_utils", + COMMON_MODULE, + HF_DIRECT_MODULE, + DISTRIBUTED_MODULE, + MODULE_PATH, + ) + saved_modules = _unit_stubs.save_sys_modules(module_names) + dist_attributes = ("get_rank", "get_world_size", "new_group", "barrier", "gather_object") + saved_dist = {name: getattr(torch_dist, name) for name in dist_attributes} + for name in module_names: + sys.modules.pop(name, None) + + _unit_stubs.install_megatron_mpu_stub() + _unit_stubs.install_ray_stub() + _unit_stubs.install_vime_distributed_utils_stub() + + iterator = MagicMock() + iterator.megatron_local_param_info_buckets = None + hf_direct = types.ModuleType(HF_DIRECT_MODULE) + hf_direct.HfWeightIteratorDirect = MagicMock(return_value=iterator) + sys.modules[HF_DIRECT_MODULE] = hf_direct + + distributed = types.ModuleType(DISTRIBUTED_MODULE) + distributed.post_process_weights = MagicMock() + sys.modules[DISTRIBUTED_MODULE] = distributed + + torch_dist.get_rank = MagicMock(return_value=0) + torch_dist.get_world_size = MagicMock(return_value=1) + torch_dist.new_group = MagicMock(return_value="slot-group") + torch_dist.barrier = MagicMock() + torch_dist.gather_object = MagicMock() + + try: + yield importlib.import_module(MODULE_PATH) + finally: + _unit_stubs.restore_sys_modules(saved_modules) + for name, value in saved_dist.items(): + setattr(torch_dist, name, value) + + +@dataclass +class RemoteCall: + args: tuple + kwargs: dict + + +class RecordingRemoteMethod: + def __init__(self): + self.calls: list[RemoteCall] = [] + + def remote(self, *args, **kwargs): + self.calls.append(RemoteCall(args, kwargs)) + return "ref" + + +@dataclass +class RecordingEngine: + init_weight_transfer_engine: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + start_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + start_draft_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + update_weights: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + finish_weight_update: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + pause_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + flush_cache: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + continue_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + + +class RecordingTrainer: + def __init__(self, client, *, source=None, fail=False): + self.client = client + self.source = source + self.fail = fail + self.draft_states = [] + self.source_draft_states = [] + self.shutdown_calls = 0 + + def send_weights(self): + self.draft_states.append(self.client.draft) + self.source_draft_states.append(getattr(self.source, "draft", False)) + self.client.start_weight_update() + if self.fail: + raise RuntimeError("transfer failed") + self.client.update_weights({"names": []}) + self.client.finish_weight_update() + + def shutdown(self): + self.shutdown_calls += 1 + + +def _args(**overrides): + values = { + "actor_num_nodes": 1, + "actor_num_gpus_per_node": 2, + "rollout_num_gpus_per_engine": 2, + "update_weight_buffer_size": 1024, + "enable_mtp_training": False, + "dspark_enabled": False, + "dspark_pretrained_model": None, + "vllm_speculative_config": None, + } + values.update(overrides) + return Namespace(**values) + + +def _updater(update_module, **overrides): + updater = object.__new__(update_module.UpdateWeightFromTensor) + updater.args = _args(**overrides) + updater.model = [] + updater.weights_getter = lambda: {} + updater.rank = 0 + updater.model_name = "test" + updater.quantization_config = None + updater.weight_version = 0 + updater.update_weight_metrics = {} + updater._hf_weight_iterator = MagicMock() + updater._full_param_info_buckets = None + updater._non_expert_param_info_buckets = None + updater._source = types.SimpleNamespace(draft=False) + updater._ipc_gather_group = None + updater._ipc_gather_src = None + updater._ipc_engine = None + updater._expert_transfer_plan = [] + updater._native_trainers = [] + updater._all_rollout_engines = [] + updater.rollout_engines = [] + return updater + + +def _install_ipc_trainer_stubs(monkeypatch, created): + factory_module = types.ModuleType("vllm.distributed.weight_transfer.factory") + ipc_module = types.ModuleType("vllm.distributed.weight_transfer.ipc_engine") + + class IPCTrainerInitInfo: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class Factory: + @staticmethod + def trainer_init(init_info, *, client, source): + created.append((init_info, client, source)) + return RecordingTrainer(client) + + factory_module.WeightTransferTrainerFactory = Factory + ipc_module.IPCTrainerInitInfo = IPCTrainerInitInfo + monkeypatch.setitem(sys.modules, factory_module.__name__, factory_module) + monkeypatch.setitem(sys.modules, ipc_module.__name__, ipc_module) + + +@pytest.mark.unit +def test_connect_uses_native_ipc_and_nccl_trainers(update_module, monkeypatch): + from vime.platforms import get_platform + + monkeypatch.setattr(update_module, "current_platform", lambda: get_platform("cuda")) + updater = _updater(update_module) + old_trainer = RecordingTrainer(object()) + updater._native_trainers = [old_trainer] + engines = [RecordingEngine(), RecordingEngine()] + ipc_created = [] + nccl_created = [] + _install_ipc_trainer_stubs(monkeypatch, ipc_created) + monkeypatch.setattr(update_module, "configure_expert_routing", lambda **kwargs: (None, [])) + + def create_nccl(client, source, gpu_counts): + nccl_created.append((client, source, gpu_counts)) + return RecordingTrainer(client) + + monkeypatch.setattr(update_module, "create_nccl_trainer", create_nccl) + updater.connect_rollout_engines( + engines, + object(), + engine_gpu_counts=[2, 2], + engine_gpu_offsets=[0, 2], + ) + + assert old_trainer.shutdown_calls == 1 + assert len(updater._native_trainers) == 2 + assert ipc_created[0][0].packed is True + assert ipc_created[0][2] is updater._source + assert nccl_created[0][1:] == (updater._source, [2]) + + +@pytest.mark.unit +def test_connect_keeps_rank_local_expert_fallback(update_module, monkeypatch): + updater = _updater(update_module) + engine = RecordingEngine() + plan = [object()] + monkeypatch.setattr(update_module, "configure_expert_routing", lambda **kwargs: ([], plan)) + monkeypatch.setattr(update_module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(update_module.dist, "new_group", lambda **kwargs: "slot-group") + + updater.connect_rollout_engines( + [engine], + object(), + engine_gpu_counts=[2], + engine_gpu_offsets=[0], + ) + updater.connect_rollout_engines( + [engine], + object(), + engine_gpu_counts=[2], + engine_gpu_offsets=[0], + ) + + assert updater._native_trainers == [] + assert updater._ipc_engine is engine + assert updater._ipc_gather_src == 0 + assert len(engine.init_weight_transfer_engine.calls) == 2 + + +@pytest.mark.unit +def test_native_update_runs_main_and_draft_lifecycles(update_module, monkeypatch): + updater = _updater( + update_module, + enable_mtp_training=True, + vllm_speculative_config={"method": "mtp"}, + ) + engine = RecordingEngine() + updater._all_rollout_engines = [engine] + client = update_module.VimeRayWeightSyncClient([engine], lambda: updater.weight_version) + trainer = RecordingTrainer(client) + updater._native_trainers = [trainer] + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + updater.update_weights() + + assert trainer.draft_states == [False, True] + assert len(engine.pause_generation.calls) == 1 + assert len(engine.start_weight_update.calls) == 1 + assert len(engine.start_draft_weight_update.calls) == 1 + assert len(engine.finish_weight_update.calls) == 2 + assert len(engine.continue_generation.calls) == 1 + + +@pytest.mark.unit +def test_native_dspark_update_uses_draft_source_and_lifecycle(update_module, monkeypatch): + updater = _updater( + update_module, + dspark_enabled=True, + vllm_speculative_config={"method": "dspark"}, + ) + engine = RecordingEngine() + updater._all_rollout_engines = [engine] + client = update_module.VimeRayWeightSyncClient([engine], lambda: updater.weight_version) + trainer = RecordingTrainer(client, source=updater._source) + updater._native_trainers = [trainer] + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + updater.update_weights() + + assert trainer.draft_states == [False, True] + assert trainer.source_draft_states == [False, True] + assert updater._source.draft is False + assert len(engine.start_draft_weight_update.calls) == 1 + + +@pytest.mark.unit +def test_failed_native_update_does_not_resume_generation(update_module, monkeypatch): + updater = _updater(update_module) + engine = RecordingEngine() + updater._all_rollout_engines = [engine] + client = update_module.VimeRayWeightSyncClient([engine], lambda: updater.weight_version) + updater._native_trainers = [RecordingTrainer(client, fail=True)] + monkeypatch.setattr(update_module.dist, "barrier", lambda *args, **kwargs: None) + + with pytest.raises(RuntimeError, match="transfer failed"): + updater.update_weights() + + assert engine.continue_generation.calls == [] + + +@pytest.mark.unit +def test_native_ipc_buffer_covers_largest_reconstructed_tensor(update_module, monkeypatch): + dense = ParamInfo("dense", torch.float16, (8,), {}, 16, 0) + expert = ParamInfo("layers.0.experts.0.weight", torch.float16, (8,), {}, 16, 0) + monkeypatch.setattr(update_module.mpu, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(update_module.mpu, "get_expert_tensor_parallel_world_size", lambda: 8) + + size = update_module._native_ipc_buffer_size(_args(update_weight_buffer_size=32), [[dense, expert]]) + + assert size == 128 + + +@pytest.mark.unit +def test_build_packed_ipc_update_info_uses_vllm_wire_format(update_module, monkeypatch): + packed_module = types.ModuleType("vllm.distributed.weight_transfer.packed_tensor") + packed_tensor = torch.zeros(12, dtype=torch.uint8) + packed_module.pack_tensors = lambda *args, **kwargs: types.SimpleNamespace( + packed_tensor=packed_tensor, + names=["a", "b"], + dtypes=[torch.float16, torch.float32], + shapes=[[2], [2]], + tensor_sizes=[4, 8], + ) + monkeypatch.setitem(sys.modules, packed_module.__name__, packed_module) + monkeypatch.setattr(torch.multiprocessing.reductions, "reduce_tensor", lambda tensor: (None, ("ipc",))) + monkeypatch.setattr(update_module.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + update_module.torch.cuda, + "get_device_properties", + lambda device: types.SimpleNamespace(uuid="GPU-0"), + ) + + update_info, weight_ref = update_module._build_packed_ipc_update_info( + [("a", torch.zeros(2, dtype=torch.float16)), ("b", torch.zeros(2))] + ) + + assert update_info == { + "names": ["a", "b"], + "dtype_names": ["float16", "float32"], + "shapes": [[2], [2]], + "tensor_sizes": [4, 8], + "ipc_handles": {"GPU-0": ("ipc",)}, + } + assert weight_ref is packed_tensor + + +@pytest.mark.unit +def test_single_worker_rank_local_payload(update_module, monkeypatch): + engine = RecordingEngine() + local_info = {"names": ["a"], "ipc_handles": {"GPU-0": ("ipc",)}} + monkeypatch.setattr(update_module, "_build_packed_ipc_update_info", lambda tensors: (local_info, "weight-ref")) + monkeypatch.setattr(update_module.dist, "get_world_size", lambda group: 1) + + refs, weight_ref = update_module._send_to_colocated_engine( + [("a", torch.zeros(1))], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group="slot-group", + ) + + assert refs == ["ref"] + assert weight_ref == "weight-ref" + assert engine.update_weights.calls[0].args == (local_info,) + + +@pytest.mark.unit +def test_multi_worker_rank_local_payload(update_module, monkeypatch): + engine = RecordingEngine() + local_info = {"names": ["a"], "ipc_handles": {"GPU-0": ("ipc-0",)}} + remote_info = {"names": ["b"], "ipc_handles": {"GPU-1": ("ipc-1",)}} + monkeypatch.setattr(update_module, "_build_packed_ipc_update_info", lambda tensors: (local_info, "weight-ref")) + monkeypatch.setattr(update_module.dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(update_module.dist, "get_rank", lambda: 0) + + def gather_object(local, object_gather_list, dst, group): + object_gather_list[:] = [local, remote_info] + + monkeypatch.setattr(update_module.dist, "gather_object", gather_object) + + refs, _ = update_module._send_to_colocated_engine( + [("a", torch.zeros(1))], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group="slot-group", + ) + + assert refs == ["ref"] + assert engine.update_weights.calls[0].args == ([local_info, remote_info],) diff --git a/tests/utils/test_vllm_arguments.py b/tests/utils/test_vllm_arguments.py new file mode 100644 index 000000000..6eeb10424 --- /dev/null +++ b/tests/utils/test_vllm_arguments.py @@ -0,0 +1,360 @@ +"""CPU unit tests for ``vime.backends.vllm_utils.arguments``.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs +import pytest + +_real_vllm = _unit_stubs.real_module_available("vllm") +requires_vllm = pytest.mark.skipif(not _real_vllm, reason="requires real vllm install") + +_unit_stubs.install_vllm_cli_stubs() + +NUM_GPUS = 0 + + +@pytest.mark.unit +@pytest.mark.parametrize("preloaded", [False, True]) +def test_real_module_available_rejects_missing_package_and_stub(monkeypatch, preloaded): + name = "_vime_missing_optional_dependency" + if preloaded: + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + assert not _unit_stubs.real_module_available(name) + + +@pytest.mark.unit +def test_real_module_available_accepts_loaded_real_module(): + assert _unit_stubs.real_module_available("sys") + + +@pytest.fixture(scope="module") +def args_mod(): + from vime.backends.vllm_utils import arguments as mod # noqa: PLC0415 + + return mod + + +def _ns(**overrides): + base = dict( + vllm_data_parallel_size=1, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + rollout_num_gpus_per_engine=4, + rollout_top_k=-1, + rollout_top_p=1.0, + vllm_router_ip=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.mark.unit +def test_validate_args_pp1(args_mod): + ns = _ns() + args_mod.validate_args(ns) + assert ns.vllm_pp_size == 1 + assert ns.vllm_dp_size == 1 + assert not hasattr(ns, "vllm_tp_size") + + +@pytest.mark.unit +def test_validate_args_records_pp_dp_but_no_global_tp(args_mod): + # validate_args records pp/dp on the namespace but must not precompute a global TP, even when + # pp>1 and dp>1. Per-engine TP = gpus_per_engine // (pp * pcp * dp) is resolved at launch time. + ns = _ns(vllm_pipeline_parallel_size=2, vllm_data_parallel_size=2) + args_mod.validate_args(ns) + assert ns.vllm_pp_size == 2 + assert ns.vllm_dp_size == 2 + assert not hasattr(ns, "vllm_tp_size") + + +@pytest.mark.unit +def test_validate_args_no_longer_raises_on_pp_indivisible(args_mod): + ns = _ns(vllm_pipeline_parallel_size=3, rollout_num_gpus_per_engine=4) + args_mod.validate_args(ns) # must not raise + assert ns.vllm_pp_size == 3 + + +@pytest.mark.unit +def test_validate_args_router_ipv6_wrapped(args_mod): + ns = _ns(vllm_router_ip="::1") + args_mod.validate_args(ns) + assert ns.vllm_router_ip == "[::1]" + + +@pytest.mark.unit +def test_validate_args_router_ipv6_already_wrapped_unchanged(args_mod): + ns = _ns(vllm_router_ip="[::1]") + args_mod.validate_args(ns) + assert ns.vllm_router_ip == "[::1]" + + +@pytest.mark.unit +def test_validate_args_router_ipv4_unchanged(args_mod): + ns = _ns(vllm_router_ip="127.0.0.1") + args_mod.validate_args(ns) + assert ns.vllm_router_ip == "127.0.0.1" + + +@pytest.mark.unit +def test_validate_args_router_none_noop(args_mod): + ns = _ns(vllm_router_ip=None) + args_mod.validate_args(ns) + assert ns.vllm_router_ip is None + + +@pytest.mark.unit +def test_validate_args_rejects_unbounded_top_p_replay(args_mod): + with pytest.raises(ValueError, match="requires --rollout-top-k > 0"): + args_mod.validate_args(_ns(rollout_top_p=0.95)) + + +@pytest.mark.unit +def test_validate_args_accepts_bounded_top_p_replay(args_mod): + args_mod.validate_args(_ns(rollout_top_p=0.95, rollout_top_k=20)) + + +@pytest.mark.unit +def test_add_vllm_router_arguments_registers_vllm_prefix(args_mod): + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_router_arguments(parser) + flags = {s for a in parser._actions for s in a.option_strings} + assert "--vllm-router-ip" in flags + assert "--vllm-router-port" in flags + assert "--vllm-router-request-timeout-secs" in flags + + +@pytest.mark.unit +def test_add_vllm_router_arguments_dests(args_mod): + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_router_arguments(parser) + dests = {a.dest for a in parser._actions if a.option_strings} + assert "vllm_router_ip" in dests + assert "vllm_router_port" in dests + assert "vllm_router_request_timeout_secs" in dests + + +@pytest.mark.unit +def test_add_vllm_router_arguments_no_unprefixed_names(args_mod): + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_router_arguments(parser) + flags = {s for a in parser._actions for s in a.option_strings} + dests = {a.dest for a in parser._actions if a.option_strings} + assert "--router-ip" not in flags + assert "--router-port" not in flags + assert "router_ip" not in dests + assert "router_port" not in dests + + +@pytest.mark.unit +def test_add_vllm_router_arguments_parses_real_values(args_mod): + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_router_arguments(parser) + parsed, _ = parser.parse_known_args( + ["--vllm-router-ip", "10.0.0.1", "--vllm-router-port", "8000", "--vllm-router-request-timeout-secs", "30"] + ) + assert parsed.vllm_router_ip == "10.0.0.1" + assert parsed.vllm_router_port == 8000 + assert parsed.vllm_router_request_timeout_secs == 30 + + +@pytest.mark.unit +def test_add_vllm_router_arguments_defaults_to_cache_aware(args_mod): + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_router_arguments(parser) + parsed, _ = parser.parse_known_args([]) + assert parsed.router_policy == "cache_aware" + + +@pytest.mark.unit +def test_add_vllm_arguments_overrides_router_balance_threshold_defaults(args_mod, monkeypatch): + _patch_device_config(monkeypatch) + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_arguments(parser) + parsed, _ = parser.parse_known_args([]) + assert parsed.router_balance_abs_threshold == 10 + assert parsed.router_balance_rel_threshold == 1.2 + + +def _patch_device_config(monkeypatch): + """Avoid device detection and third-party plugin loading in parser tests.""" + try: + from vllm.config.device import DeviceConfig + from vllm.engine import arg_utils + + monkeypatch.setattr(DeviceConfig, "__post_init__", lambda self: setattr(self, "device_type", "cpu")) + monkeypatch.setattr(arg_utils, "load_general_plugins", lambda: None) + except ImportError: + pass + + +@pytest.mark.unit +@requires_vllm +def test_add_vllm_arguments_prefixes_regular_engine_flags(args_mod, monkeypatch): + _patch_device_config(monkeypatch) + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_arguments(parser) + flags = {s for a in parser._actions for s in a.option_strings} + assert "--vllm-server-concurrency" in flags + assert "--vllm-tool-call-parser" in flags + assert "--vllm-weight-sync-packed" not in flags + assert "--no-vllm-weight-sync-packed" not in flags + + +@pytest.mark.unit +def test_add_vllm_arguments_skips_orchestrator_owned_fields(args_mod, monkeypatch): + _patch_device_config(monkeypatch) + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_arguments(parser) + flags = {s for a in parser._actions for s in a.option_strings} + dests = {a.dest for a in parser._actions if a.option_strings} + assert "--vllm-seed" not in flags + assert "--vllm-host" not in flags + assert "--vllm-master-addr" not in flags + assert "--vllm-tensor-parallel-size" not in flags + assert "seed" not in dests + assert "host" not in dests + assert "master_addr" not in dests + assert "tensor_parallel_size" not in dests + + +@pytest.mark.unit +@requires_vllm +def test_add_vllm_arguments_parses_prefixed_engine_values(args_mod, monkeypatch): + _patch_device_config(monkeypatch) + parser = argparse.ArgumentParser(add_help=False) + args_mod.add_vllm_arguments(parser) + parsed, _ = parser.parse_known_args(["--vllm-server-concurrency", "128", "--vllm-tool-call-parser", "qwen3_coder"]) + assert parsed.vllm_server_concurrency == 128 + assert parsed.vllm_tool_call_parser == "qwen3_coder" + + +@pytest.mark.unit +def test_parse_args_tp_default_no_pp(args_mod, monkeypatch): + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr(sys, "argv", ["train.py", "--rollout-num-gpus-per-engine", "4"]) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 4 + + +@pytest.mark.unit +def test_parse_args_tp_default_with_pp(args_mod, monkeypatch): + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr( + sys, + "argv", + ["train.py", "--rollout-num-gpus-per-engine", "4", "--vllm-pipeline-parallel-size", "2"], + ) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 2 + + +@pytest.mark.unit +def test_parse_args_default_attribute_set_even_without_register(args_mod, monkeypatch): + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr(sys, "argv", ["train.py", "--rollout-num-gpus-per-engine", "8"]) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 8 + + +@pytest.mark.unit +def test_parse_args_tp_default_with_dp(args_mod, monkeypatch): + """TP auto-compute must divide by DP: TP = gpus / (PP * PCP * DP).""" + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr( + sys, + "argv", + ["train.py", "--rollout-num-gpus-per-engine", "8", "--vllm-data-parallel-size", "4"], + ) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 2 # 8 / (1 * 4) = 2 + + +@pytest.mark.unit +def test_parse_args_tp_default_with_pcp_and_dp(args_mod, monkeypatch): + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr( + sys, + "argv", + [ + "train.py", + "--rollout-num-gpus-per-engine", + "8", + "--vllm-prefill-context-parallel-size", + "2", + "--vllm-data-parallel-size", + "2", + ], + ) + + ns = args_mod.vllm_parse_args() + + assert ns.vllm_tensor_parallel_size == 2 + + +@pytest.mark.unit +def test_parse_args_tp_default_with_pp_and_dp(args_mod, monkeypatch): + """TP = gpus / (PP * DP) when both PP and DP are set.""" + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr( + sys, + "argv", + [ + "train.py", + "--rollout-num-gpus-per-engine", + "8", + "--vllm-pipeline-parallel-size", + "2", + "--vllm-data-parallel-size", + "2", + ], + ) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 2 # 8 / (2 * 2) = 2 + + +@pytest.mark.unit +def test_parse_args_tp_default_dp1_unchanged(args_mod, monkeypatch): + """DP=1 (default) must not change existing TP behavior.""" + monkeypatch.setattr(args_mod, "add_vllm_arguments", lambda p: p) + monkeypatch.setattr( + sys, + "argv", + ["train.py", "--rollout-num-gpus-per-engine", "4", "--vllm-data-parallel-size", "1"], + ) + ns = args_mod.vllm_parse_args() + assert ns.vllm_tensor_parallel_size == 4 # 4 / (1 * 1) = 4 + + +@pytest.mark.unit +def test_validate_args_rejects_conflicting_rollout_external_and_vllm_config(args_mod): + ns = _ns(rollout_external=True, vllm_config="/tmp/vllm.yaml") + with pytest.raises(AssertionError, match="vllm_config cannot be set"): + args_mod.validate_args(ns) + + +@pytest.mark.unit +def test_validate_args_rejects_conflicting_prefill_and_vllm_config(args_mod): + ns = _ns(prefill_num_servers=2, vllm_config="/tmp/vllm.yaml") + with pytest.raises(AssertionError, match="mutually exclusive"): + args_mod.validate_args(ns) + + +@pytest.mark.unit +def test_validate_args_rejects_prefill_and_rollout_external(args_mod): + ns = _ns(prefill_num_servers=2, rollout_external=True) + with pytest.raises(AssertionError, match="cannot be set"): + args_mod.validate_args(ns) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_vllm_config.py b/tests/utils/test_vllm_config.py index 42e873e18..4f5556b57 100644 --- a/tests/utils/test_vllm_config.py +++ b/tests/utils/test_vllm_config.py @@ -2,18 +2,17 @@ import sys import tempfile +from argparse import Namespace from pathlib import Path import pytest import yaml -_tests_root = Path(__file__).resolve().parents[1] -if str(_tests_root) not in sys.path: - sys.path.insert(0, str(_tests_root)) +NUM_GPUS = 0 -import _unit_stubs - -_unit_stubs.install_rollout_optional_stubs() +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) def _write_yaml(data: dict) -> str: @@ -40,6 +39,8 @@ def test_update_weights_defaults_to_none(self): ) config = VllmConfig.from_yaml(path) assert len(config.models) == 1 + # Parsed default is None; VllmConfig.resolve() later infers True/False from + # whether model_path matches args.hf_checkpoint. assert config.models[0].update_weights is None def test_update_weights_explicit_false(self): @@ -93,6 +94,395 @@ def test_multi_model_total_gpus(self): config = VllmConfig.from_yaml(path) assert config.total_num_gpus == 12 + def test_config_allows_model_with_no_server_groups(self): + """A model with no server groups can expose a router without local engines.""" + from vime.backends.vllm_utils.vllm_config import VllmConfig + + path = _write_yaml({"vllm": [{"name": "default", "server_groups": []}]}) + + config = VllmConfig.from_yaml(path) + + assert len(config.models) == 1 + assert config.models[0].name == "default" + assert config.models[0].server_groups == [] + assert config.total_num_gpus == 0 + + +class TestZeroGpuRolloutConfig: + def test_resolve_default_zero_gpu_config_has_no_server_groups(self): + from vime.backends.vllm_utils.vllm_config import resolve_vllm_config + + args = Namespace(vllm_config=None, prefill_num_servers=None, rollout_num_gpus=0) + + config = resolve_vllm_config(args) + + assert len(config.models) == 1 + assert config.models[0].name == "default" + assert config.models[0].server_groups == [] + assert config.total_num_gpus == 0 + + def test_zero_gpu_config_takes_precedence_over_prefill_num_servers(self): + from vime.backends.vllm_utils.vllm_config import resolve_vllm_config + + args = Namespace(vllm_config=None, prefill_num_servers=1, rollout_num_gpus=0) + + config = resolve_vllm_config(args) + + assert config.models[0].server_groups == [] + assert config.total_num_gpus == 0 + + def test_start_rollout_servers_zero_gpu_starts_router_without_engines(self, monkeypatch): + from vime.backends.vllm_utils import deployment + + def fake_start_router(args, *, has_pd_disaggregation=False, force_new=False): + assert has_pd_disaggregation is False + assert force_new is False + return "127.0.0.1", 3456, None + + monkeypatch.setattr(deployment, "_start_router", fake_start_router) + args = Namespace( + rollout_external=False, + vllm_config=None, + prefill_num_servers=None, + rollout_num_gpus=0, + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + debug_train_only=False, + debug_rollout_only=False, + colocate=False, + actor_num_nodes=1, + actor_num_gpus_per_node=8, + offload_rollout=False, + hf_checkpoint="/tmp/hf", + ) + + servers, init_handles = deployment.start_rollout_servers(args, pg=(None, [], [])) + + assert list(servers) == ["default"] + assert init_handles == [] + server = servers["default"] + assert server.router_ip == "127.0.0.1" + assert server.router_port == 3456 + assert server.server_groups == [] + assert server.engines == [] + assert args.vllm_router_ip == "127.0.0.1" + assert args.vllm_router_port == 3456 + assert args.vllm_model_routers == {"default": ("127.0.0.1", 3456)} + + def test_server_group_parallel_config_derives_tp_from_overridden_pp(self): + from vime.backends.vllm_utils.engine_group import ServerGroup + + args = Namespace( + num_gpus_per_node=8, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + vllm_data_parallel_size=1, + vllm_dp_size=1, + vllm_enable_expert_parallel=True, + ) + + group = ServerGroup( + args=args, + pg=None, + all_engines=[object()], + num_gpus_per_engine=32, + num_new_engines=1, + vllm_overrides={"pipeline_parallel_size": 2}, + ) + + assert group.parallel_config() == { + "tp_size": 16, + "pp_size": 2, + "pcp_size": 1, + "dp_size": 1, + "enable_expert_parallel": True, + "ep_size": 16, + } + + def test_server_group_parallel_config_derives_tp_from_overridden_pcp(self): + from vime.backends.vllm_utils.engine_group import ServerGroup + + args = Namespace( + num_gpus_per_node=8, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + vllm_data_parallel_size=1, + vllm_dp_size=1, + vllm_enable_expert_parallel=True, + ) + group = ServerGroup( + args=args, + pg=None, + all_engines=[object()], + num_gpus_per_engine=8, + num_new_engines=1, + vllm_overrides={"prefill_context_parallel_size": 2, "data_parallel_size": 2}, + ) + + assert group.parallel_config() == { + "tp_size": 2, + "pp_size": 1, + "pcp_size": 2, + "dp_size": 2, + "enable_expert_parallel": True, + "ep_size": 8, + } + + def test_vllm_server_args_derive_tp_from_overridden_pp(self, monkeypatch): + from vime.backends.vllm_utils import vllm_engine + + monkeypatch.setattr(vllm_engine, "_VLLM_SERVER_FIELDS", frozenset()) + + args = Namespace( + hf_checkpoint="/tmp/hf", + seed=1, + offload_rollout=False, + rollout_num_gpus_per_engine=32, + num_gpus_per_node=8, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + vllm_data_parallel_size=1, + vllm_dp_size=1, + vllm_enable_expert_parallel=True, + use_rollout_routing_replay=False, + fp16=False, + colocate=False, + rollout_max_context_len=None, + vllm_max_model_len=None, + ) + + kwargs, _ = vllm_engine._compute_server_args( + args, + rank=0, + dist_init_addr="127.0.0.1:12345", + host="127.0.0.1", + port=30000, + base_gpu_id=0, + vllm_overrides={"pipeline_parallel_size": 2}, + num_gpus_per_engine=32, + ) + + assert kwargs["pipeline_parallel_size"] == 2 + assert kwargs["tensor_parallel_size"] == 16 + + def test_offload_rollout_enables_vllm_sleep_mode(self, monkeypatch): + from vime.backends.vllm_utils import vllm_engine + + monkeypatch.setattr(vllm_engine, "_VLLM_SERVER_FIELDS", frozenset()) + + args = Namespace( + hf_checkpoint="/tmp/hf", + seed=1, + offload_rollout=True, + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + vllm_data_parallel_size=1, + vllm_dp_size=1, + vllm_enable_expert_parallel=False, + use_rollout_routing_replay=False, + fp16=False, + colocate=False, + rollout_max_context_len=None, + vllm_max_model_len=None, + vllm_enable_sleep_mode=False, + ) + compute_kwargs = { + "rank": 0, + "dist_init_addr": "127.0.0.1:12345", + "host": "127.0.0.1", + "port": 30000, + "base_gpu_id": 0, + } + + kwargs, _ = vllm_engine._compute_server_args(args, **compute_kwargs) + assert kwargs["enable_sleep_mode"] is True + assert args.vllm_enable_sleep_mode is True + + def test_start_rollout_servers_defers_engine_wait(self, monkeypatch): + from vime.backends.vllm_utils import deployment, disaggregation + from vime.backends.vllm_utils.engine_group import ServerGroup + + def fake_start_router(args, *, has_pd_disaggregation=False, force_new=False): + assert has_pd_disaggregation is False + assert force_new is False + return "127.0.0.1", 3456, None + + def fake_start_engines(self, port_cursors=None): + self.all_engines = [object() for _ in self.all_engines] + return [f"init-{self.rank_offset}"], port_cursors or {} + + def fail_if_waited(_refs): + pytest.fail("regular deployment must not wait for engine initialization") + + monkeypatch.setattr(deployment, "_start_router", fake_start_router) + monkeypatch.setattr(ServerGroup, "start_engines", fake_start_engines) + monkeypatch.setattr(disaggregation.ray, "get", fail_if_waited) + + args = Namespace( + rollout_external=False, + vllm_config=None, + prefill_num_servers=None, + rollout_num_gpus=2, + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + debug_train_only=False, + debug_rollout_only=False, + colocate=False, + actor_num_nodes=1, + actor_num_gpus_per_node=8, + offload_rollout=False, + hf_checkpoint="/tmp/hf", + ) + + servers, init_handles = deployment.start_rollout_servers(args, pg=(None, [], [])) + + assert list(servers) == ["default"] + assert init_handles == ["init-0"] + + def test_start_rollout_servers_routes_pd_to_disaggregated_deployment(self, monkeypatch): + from vime.backends.vllm_utils import deployment + from vime.backends.vllm_utils.engine_group import ServerGroup + from vime.backends.vllm_utils.vllm_config import ModelConfig, ServerGroupConfig, VllmConfig + + def fake_start_router( + args, + *, + has_pd_disaggregation=False, + force_new=False, + bind=None, + prefill_urls=None, + decode_urls=None, + ): + assert has_pd_disaggregation is True + assert force_new is False + assert bind is not None + assert prefill_urls == [("http://prefill", None)] + assert decode_urls == ["http://decode"] + return "127.0.0.1", 3456, None + + def fake_resolve_vllm_config(args): + return VllmConfig( + models=[ + ModelConfig( + name="default", + server_groups=[ + ServerGroupConfig(worker_type="prefill", num_gpus=1), + ServerGroupConfig(worker_type="decode", num_gpus=1), + ], + ) + ] + ) + + def fake_start_engines(self, port_cursors=None): + self.all_engines = [object() for _ in self.all_engines] + return [f"{self.worker_type}-init-{self.rank_offset}"], port_cursors or {} + + monkeypatch.setattr(deployment, "_start_router", fake_start_router) + monkeypatch.setattr(deployment, "resolve_vllm_config", fake_resolve_vllm_config) + monkeypatch.setattr( + deployment, + "collect_pd_urls", + lambda _groups: ([("http://prefill", None)], ["http://decode"]), + ) + monkeypatch.setattr(ServerGroup, "start_engines", fake_start_engines) + + args = Namespace( + rollout_external=False, + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + debug_train_only=False, + debug_rollout_only=False, + colocate=False, + actor_num_nodes=1, + actor_num_gpus_per_node=8, + offload_rollout=False, + hf_checkpoint="/tmp/hf", + ) + + servers, init_handles = deployment.start_rollout_servers(args, pg=(None, [], [])) + + groups = servers["default"].server_groups + assert [group.worker_type for group in groups] == ["prefill", "decode"] + assert init_handles == ["prefill-init-0", "decode-init-1"] + + def test_start_rollout_servers_waits_for_epd_encoder_before_non_encoder(self, monkeypatch): + from vime.backends.vllm_utils import deployment, disaggregation + from vime.backends.vllm_utils.engine_group import ServerGroup + from vime.backends.vllm_utils.vllm_config import ModelConfig, ServerGroupConfig, VllmConfig + + class FakeRemoteMethod: + def __init__(self, value): + self.value = value + + def remote(self): + return self.value + + class FakeEngine: + def __init__(self, url_ref): + self.get_url = FakeRemoteMethod(url_ref) + + def fake_start_router(args, *, has_pd_disaggregation=False, force_new=False): + assert has_pd_disaggregation is False + assert force_new is False + return "127.0.0.1", 3456, None + + def fake_resolve_vllm_config(args): + return VllmConfig( + models=[ + ModelConfig( + name="default", + server_groups=[ + ServerGroupConfig(worker_type="encoder", num_gpus=1), + ServerGroupConfig(worker_type="regular", num_gpus=1), + ], + ) + ] + ) + + def fake_start_engines(self, port_cursors=None): + if self.worker_type == "encoder": + self.all_engines = [FakeEngine("encoder-url-ref") for _ in self.all_engines] + else: + self.all_engines = [object() for _ in self.all_engines] + return [f"{self.worker_type}-init-{self.rank_offset}"], port_cursors or {} + + ray_get_calls = [] + + def fake_ray_get(refs): + ray_get_calls.append(refs) + if refs == ["encoder-url-ref"]: + return ["http://encoder"] + return None + + monkeypatch.setattr(deployment, "_start_router", fake_start_router) + monkeypatch.setattr(deployment, "resolve_vllm_config", fake_resolve_vllm_config) + monkeypatch.setattr(ServerGroup, "start_engines", fake_start_engines) + monkeypatch.setattr(disaggregation.ray, "get", fake_ray_get) + + args = Namespace( + rollout_external=False, + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + debug_train_only=False, + debug_rollout_only=False, + colocate=False, + actor_num_nodes=1, + actor_num_gpus_per_node=8, + offload_rollout=False, + hf_checkpoint="/tmp/hf", + ) + + servers, init_handles = deployment.start_rollout_servers(args, pg=(None, [], [])) + + groups = servers["default"].server_groups + assert [group.worker_type for group in groups] == ["encoder", "regular"] + assert groups[1].vllm_overrides["language_only"] is True + assert groups[1].vllm_overrides["encoder_urls"] == ["http://encoder"] + assert init_handles == ["regular-init-1"] + assert ray_get_calls == [["encoder-init-0"], ["encoder-url-ref"]] + class TestGetModelUrl: def test_get_model_url_basic(self): @@ -140,4 +530,4 @@ def test_get_model_url_no_routers(self): if __name__ == "__main__": - pytest.main([__file__, "-v"]) + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py new file mode 100644 index 000000000..bd540b0a5 --- /dev/null +++ b/tests/utils/test_vllm_engine.py @@ -0,0 +1,980 @@ +"""CPU unit tests for ``vime.backends.vllm_utils.vllm_engine``.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs +import pytest +import requests +import torch + +_unit_stubs.install_vllm_cli_stubs() + +from vime.backends.vllm_utils import vllm_engine as mod + +NUM_GPUS = 0 + + +@pytest.fixture(autouse=True) +def _patch_vllm_server_fields(monkeypatch): + """Stub _VLLM_SERVER_FIELDS so tests don't need a real vllm install.""" + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset()) + + +@pytest.fixture +def vllm_args() -> SimpleNamespace: + return SimpleNamespace( + rollout_external=True, + hf_checkpoint="/tmp/model", + vllm_router_ip=None, + vllm_router_port=None, + num_gpus_per_node=8, + rollout_num_gpus_per_engine=4, + colocate=False, + debug_rollout_only=False, + actor_num_gpus_per_node=4, + actor_num_nodes=1, + use_critic=False, + critic_num_gpus_per_node=0, + critic_num_nodes=0, + seed=1234, + fp16=False, + offload_rollout=False, + use_rollout_routing_replay=False, + rollout_top_p=1.0, + vllm_pipeline_parallel_size=1, + vllm_prefill_context_parallel_size=1, + vllm_data_parallel_size=1, + vllm_dp_size=1, + vllm_enable_expert_parallel=False, + ) + + +@pytest.fixture +def vllm_engine(vllm_args): + from vime.backends.vllm_utils.vllm_engine import VLLMEngine + + engine = VLLMEngine(vllm_args, rank=0) + engine.node_rank = 0 + engine.server_host = "127.0.0.1" + engine.server_port = 8765 + return engine + + +class _MockResponse: + def __init__(self, *, json_data: dict | None = None, text: str = "", status_code: int = 200): + self._json_data = json_data + self.text = text + self.status_code = status_code + # Model requests.Response.content (raw body bytes) so _response_json's empty-body + # handling (empty 200 -> {"ok": True}) is actually exercised. A JSON body is non-empty; + # text-only/empty bodies use the given text (b"" when empty). + self.content = json.dumps(json_data).encode() if json_data is not None else text.encode() + + def raise_for_status(self) -> None: + if self.status_code >= 400: + error = requests.exceptions.HTTPError(f"HTTP {self.status_code}") + error.response = self # type: ignore[assignment] + raise error + + def json(self) -> dict: + if self._json_data is None: + raise ValueError("no json") + return self._json_data + + +@pytest.mark.unit +def test_normalize_vllm_wake_tags_drops_unsupported(): + assert mod._normalize_vllm_wake_tags(["weights", "cuda_graph", "kv_cache"]) == ["weights", "kv_cache"] + + +@pytest.mark.unit +def test_normalize_vllm_wake_tags_empty_becomes_none(): + assert mod._normalize_vllm_wake_tags(["cuda_graph"]) is None + + +@pytest.mark.unit +def test_launch_config_single_node(vllm_args): + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 4 + vllm_args.vllm_pipeline_parallel_size = 1 + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert sa["nnodes"] == 1 + assert sa["node_rank"] == 0 + assert sa["_tp_size"] == 4 + assert sa["_pp_size"] == 1 + assert sa["_pcp_size"] == 1 + assert sa["_dp_size"] == 1 + + +@pytest.mark.unit +def test_compute_server_args_preserves_non_topology_vllm_flags(vllm_args, monkeypatch): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"server_concurrency", "tool_call_parser"})) + vllm_args.vllm_server_concurrency = 256 + vllm_args.vllm_tool_call_parser = "qwen3_coder" + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert sa["server_concurrency"] == 256 + assert sa["tool_call_parser"] == "qwen3_coder" + + +@pytest.mark.unit +def test_compute_server_args_ignores_unrecognized_vllm_attrs(vllm_args, monkeypatch): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"server_concurrency"})) + vllm_args.vllm_nonexistent_flag = "keep-out" + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert "nonexistent_flag" not in sa + + +@pytest.mark.unit +def test_launch_config_multi_node_ranks(vllm_args): + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 16 + vllm_args.vllm_pipeline_parallel_size = 2 + sa0, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr="10.0.0.1:15000", host="10.0.0.1", port=8000) + sa1, _ = mod._compute_server_args(vllm_args, rank=1, dist_init_addr="10.0.0.1:15000", host="10.0.0.2", port=8000) + assert sa0["nnodes"] == 2 + assert sa0["node_rank"] == 0 + assert sa1["node_rank"] == 1 + + +@pytest.mark.unit +def test_distributed_flags_only_when_multi_node(vllm_args): + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 4 + vllm_args.vllm_pipeline_parallel_size = 1 + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert "master_addr" not in sa + + vllm_args.rollout_num_gpus_per_engine = 16 + vllm_args.vllm_pipeline_parallel_size = 2 + sa_multi, _ = mod._compute_server_args( + vllm_args, rank=1, dist_init_addr="10.0.0.2:16000", host="10.0.0.2", port=8000 + ) + assert sa_multi["nnodes"] == 2 + assert sa_multi["node_rank"] == 1 + assert sa_multi["master_addr"] == "10.0.0.2" + assert sa_multi.get("headless") is True + assert sa_multi.get("data_parallel_backend") == "mp" + assert sa_multi.get("distributed_executor_backend") == "mp" + + +@pytest.mark.unit +def test_compute_server_args_applies_worker_type_and_bootstrap_port(vllm_args): + sa_prefill, _ = mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + worker_type="prefill", + disaggregation_bootstrap_port=12345, + ) + assert sa_prefill["kv_transfer_config"] == { + "kv_connector": "NixlConnector", + "kv_role": "kv_producer", + } + + sa_decode, _ = mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + worker_type="decode", + ) + assert sa_decode["kv_transfer_config"] == { + "kv_connector": "NixlConnector", + "kv_role": "kv_consumer", + } + + +@pytest.mark.unit +def test_compute_server_args_allows_mooncake_group_override(vllm_args): + config = { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": {"device_name": "mlx5_0,mlx5_1"}, + } + server_args, _ = mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + worker_type="prefill", + disaggregation_bootstrap_port=12345, + vllm_overrides={"kv_transfer_config": config}, + ) + assert server_args["kv_transfer_config"] == config + + +@pytest.mark.unit +def test_compute_server_args_prefill_requires_bootstrap_port(vllm_args): + with pytest.raises(AssertionError, match="disaggregation_bootstrap_port"): + mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + worker_type="prefill", + ) + + +@pytest.mark.unit +def test_compute_server_args_applies_rollout_and_dtype_flags(vllm_args): + vllm_args.use_rollout_routing_replay = True + vllm_args.vllm_speculative_config = {"method": "mtp"} + vllm_args.rollout_top_p = 0.9 + vllm_args.fp16 = True + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert sa["enable_return_routed_experts"] is True + assert sa["per_request_spec_decode_metrics"] == "summary" + assert sa["return_sampling_mask"] is True + assert sa["dtype"] == "float16" + + +@pytest.mark.unit +def test_compute_server_args_applies_max_model_len_from_rollout_context(vllm_args): + vllm_args.rollout_max_context_len = 8192 + vllm_args.vllm_max_model_len = None + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert sa["max_model_len"] == 8192 + + +@pytest.mark.unit +def test_compute_server_args_model_path_override_wins(vllm_args, monkeypatch): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"model", "server_concurrency"})) + sa, _ = mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + vllm_overrides={"model_path": "/tmp/override", "server-concurrency": 123}, + ) + assert sa["model"] == "/tmp/override" + assert sa["server_concurrency"] == 123 + + +@pytest.mark.unit +def test_compute_server_args_external_check_fields_skip_orchestration_fields(vllm_args): + sa, check_fields = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert "model" not in check_fields + assert "host" not in check_fields + assert "port" not in check_fields + assert "nnodes" in check_fields + assert "node_rank" in check_fields + assert "weight_transfer_config" in check_fields + assert sa["weight_transfer_config"] == {"backend": "nccl"} + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_colocate(vllm_args, monkeypatch): + vllm_args.colocate = True + monkeypatch.delenv("PYTHONPATH", raising=False) + env = mod._build_subprocess_env( + { + "_args": vllm_args, + "_visible_devices": "0,1", + } + ) + assert "VLLM_ALLOW_INSECURE_SERIALIZATION" in env + assert env["VLLM_ALLOW_INSECURE_SERIALIZATION"] == "1" + assert "PYTHONPATH" in env + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_drops_trainer_allocator_config(vllm_args, monkeypatch): + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + monkeypatch.setenv("PYTORCH_ALLOC_CONF", "expandable_segments:True") + + env = mod._build_subprocess_env({"_args": vllm_args, "_visible_devices": "0"}) + + assert "PYTORCH_CUDA_ALLOC_CONF" not in env + assert "PYTORCH_ALLOC_CONF" not in env + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_sets_batch_invariant_when_deterministic(vllm_args, monkeypatch): + monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False) + vllm_args.vllm_enable_deterministic_inference = True + env = mod._build_subprocess_env({"_args": vllm_args, "_visible_devices": "0"}) + assert env["VLLM_BATCH_INVARIANT"] == "1" + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_enables_v2_runner_by_default(vllm_args): + env = mod._build_subprocess_env({"_args": vllm_args, "_visible_devices": "0"}) + assert env["VLLM_USE_V2_MODEL_RUNNER"] == "1" + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_no_batch_invariant_by_default(vllm_args, monkeypatch): + monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False) + vllm_args.vllm_enable_deterministic_inference = False + env = mod._build_subprocess_env({"_args": vllm_args, "_visible_devices": "0"}) + assert "VLLM_BATCH_INVARIANT" not in env + + +@pytest.mark.unit +def test_build_vllm_subprocess_env_sets_disaggregation_side_channel(vllm_args): + env = mod._build_subprocess_env( + { + "_args": vllm_args, + "_visible_devices": "0", + "_worker_type": "prefill", + "_disaggregation_bootstrap_port": 29999, + "node_rank": 0, + "host": "10.0.0.8", + } + ) + assert env["VLLM_NIXL_SIDE_CHANNEL_HOST"] == "10.0.0.8" + assert env["VLLM_NIXL_SIDE_CHANNEL_PORT"] == "29999" + + +@pytest.mark.unit +def test_compute_server_args_adds_sleep_mode_for_offload_rollout(vllm_args): + vllm_args.offload_rollout = True + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert sa.get("enable_sleep_mode") is True + assert vllm_args.vllm_enable_sleep_mode is True + + +@pytest.mark.unit +@pytest.mark.parametrize( + "platform_name,colocate,backend", + [("cuda", False, "nccl"), ("cuda", True, "ipc"), ("npu", False, "hccl"), ("npu", True, "npu_ipc")], +) +def test_compute_server_args_uses_native_platform_backend(vllm_args, monkeypatch, platform_name, colocate, backend): + from vime.platforms import get_platform + + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"worker_extension_cls"})) + vllm_args.vllm_worker_extension_cls = "" + vllm_args.colocate = colocate + monkeypatch.setattr(mod, "current_platform", lambda: get_platform(platform_name)) + + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + + assert not sa.get("worker_extension_cls") + assert sa["weight_transfer_config"] == {"backend": backend} + + +@pytest.mark.unit +def test_compute_server_args_keeps_user_worker_extension(vllm_args, monkeypatch): + from vime.platforms import get_platform + + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"worker_extension_cls"})) + vllm_args.vllm_worker_extension_cls = "example.UserWorkerExtension" + monkeypatch.setattr(mod, "current_platform", lambda: get_platform("npu")) + + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + + assert sa["worker_extension_cls"] == "example.UserWorkerExtension" + + +@pytest.mark.unit +def test_compute_server_args_no_sleep_mode_from_colocate(vllm_args): + vllm_args.colocate = True + vllm_args.offload_rollout = False + sa, _ = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + assert "enable_sleep_mode" not in sa + assert not getattr(vllm_args, "vllm_enable_sleep_mode", False) + + +@pytest.mark.unit +def test_get_base_gpu_id_colocate(vllm_args): + vllm_args.colocate = True + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 4 + assert mod.get_base_gpu_id(vllm_args, rank=1) == 4 + + +@pytest.mark.unit +def test_start_weight_update_posts_four_phase_endpoint(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + result = vllm_engine.start_weight_update() + + assert result == {"ok": True} + assert len(calls) == 1 + assert calls[0][0] == "start_weight_update" + assert calls[0][1] == {} + + +@pytest.mark.unit +def test_start_draft_weight_update_posts_empty_body(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + result = vllm_engine.start_draft_weight_update() + + assert result == {"ok": True} + assert calls == [("start_draft_weight_update", {})] + + +@pytest.mark.unit +def test_finish_weight_update_posts_empty_body(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"done": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + result = vllm_engine.finish_weight_update() + + assert result == {"done": True} + assert calls == [("finish_weight_update", {})] + + +@pytest.mark.unit +def test_finish_weight_update_commits_version(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"done": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + assert vllm_engine.finish_weight_update(weight_version="42") == {"done": True} + assert calls == [("finish_weight_update", {"weight_version": "42"})] + assert vllm_engine._weight_version == "42" + + +@pytest.mark.unit +def test_update_weights_posts_ipc_payload(vllm_engine, monkeypatch): + posted: list[tuple[str, dict]] = [] + + def fake_post(endpoint: str, payload: dict): + posted.append((endpoint, payload)) + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + assert vllm_engine._weight_version is None + + vllm_engine.update_weights( + { + "names": ["a", "b"], + "dtype_names": ["bfloat16", "float32"], + "shapes": [[2], [1]], + "ipc_handles": {"uuid-gpu0": ("rebuild_fn", (1, 2, 3))}, + "tensor_sizes": [4, 4], + } + ) + + assert posted[0][0] == "update_weights" + sent = posted[0][1]["update_info"] + # ipc_handles are pickled for native vLLM/vLLM-Ascend parse_update_info. + # ipc_handles got cloudpickle'd into ipc_handles_pickled + assert "ipc_handles" not in sent + assert isinstance(sent["ipc_handles_pickled"], str) + assert sent["names"] == ["a", "b"] + assert sent["shapes"] == [[2], [1]] + assert sent["tensor_sizes"] == [4, 4] + assert "packed" not in sent + assert vllm_engine._weight_version is None + + +@pytest.mark.unit +def test_update_weights_serializes_rank_local_payloads(vllm_engine, monkeypatch): + posted: list[tuple[str, dict]] = [] + + def fake_post(endpoint: str, payload: dict): + posted.append((endpoint, payload)) + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + first = { + "names": ["experts.0.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[2]], + "ipc_handles": {"uuid-gpu0": ("first", ())}, + "tensor_sizes": [4], + } + third = { + **first, + "names": ["experts.2.weight"], + "ipc_handles": {"uuid-gpu2": ("third", ())}, + } + vllm_engine.update_weights([first, None, third]) + + sent = posted[0][1]["update_info"] + assert sent[1] is None + assert sent[0]["names"] == ["experts.0.weight"] + assert sent[2]["names"] == ["experts.2.weight"] + assert "ipc_handles" not in sent[0] + assert "ipc_handles" not in sent[2] + assert isinstance(sent[0]["ipc_handles_pickled"], str) + assert isinstance(sent[2]["ipc_handles_pickled"], str) + + +@pytest.mark.unit +def test_update_weights_does_not_advance_version_on_failure(vllm_engine, monkeypatch): + """Chunk POST failure must not modify the engine's committed version cache.""" + + def fake_post_fail(endpoint: str, payload: dict) -> dict: + raise RuntimeError("simulated POST failure") + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post_fail) + + vllm_engine._weight_version = "old" + with pytest.raises(RuntimeError, match="simulated POST failure"): + vllm_engine.update_weights( + {"names": [], "dtype_names": [], "shapes": [], "ipc_handles": {}, "tensor_sizes": []} + ) + assert vllm_engine._weight_version == "old" + + +@pytest.mark.unit +def test_get_weight_version_reads_vllm_weight_info(vllm_engine, monkeypatch): + monkeypatch.setattr( + mod.requests, + "get", + lambda *args, **kwargs: _MockResponse(json_data={"weight_version": "7"}), + ) + + assert vllm_engine.get_weight_version() == "7" + assert vllm_engine._weight_version == "7" + + +@pytest.mark.unit +def test_get_weight_version_preserves_uninitialized_none(vllm_engine, monkeypatch): + monkeypatch.setattr( + mod.requests, + "get", + lambda *args, **kwargs: _MockResponse(json_data={"weight_version": None}), + ) + + assert vllm_engine.get_weight_version() is None + assert vllm_engine._weight_version is None + + +@pytest.mark.unit +def test_set_weight_version_updates_vllm_and_local_cache(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"success": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + assert vllm_engine.set_weight_version("9") == {"success": True} + assert calls == [("update_weight_version", {"new_version": "9"})] + assert vllm_engine._weight_version == "9" + + +@pytest.mark.unit +def test_get_weight_version_worker_rank_returns_none_without_raise(vllm_engine): + """Worker ranks short-circuit (matches the class-wide idiom).""" + vllm_engine.node_rank = 1 + vllm_engine._weight_version = None + assert vllm_engine.get_weight_version() is None + + +@pytest.mark.unit +def test_update_weights_from_distributed_posts_update_weights_without_checkpoint_flag(vllm_engine, monkeypatch): + calls: list[dict] = [] + + def fake_make_request(endpoint: str, payload: dict) -> dict: + calls.append(payload.get("update_info", payload)) + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_make_request) + + names = ["layer.0.weight"] + dtypes = [torch.float32] + shapes = [torch.Size([2, 2])] + + vllm_engine.update_weights_from_distributed( + names, + dtypes, + shapes, + weight_version="7", + ) + + assert len(calls) == 1 + info = calls[0] + assert info["names"] == names + assert info["dtype_names"] == ["float32"] + assert info["shapes"] == [[2, 2]] + assert info["packed"] is True + assert "is_checkpoint_format" not in info + assert vllm_engine._weight_version is None + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("host", "expected_host"), + [ + ("127.0.0.1", "127.0.0.1"), + ("2001:db8::1", "[2001:db8::1]"), + ("[2001:db8::1]", "[2001:db8::1]"), + ], +) +def test_init_formats_server_and_router_hosts_for_urls(vllm_engine, monkeypatch, host, expected_host): + server_args = {} + monkeypatch.setattr(vllm_engine, "_init_external", lambda args, **kwargs: server_args.update(args)) + + vllm_engine.init( + dist_init_addr="127.0.0.1:29500", + port=8765, + nccl_port=None, + host=host, + router_ip=host, + router_port=30000, + ) + + assert server_args["host"] == expected_host + assert vllm_engine.server_host == expected_host + assert vllm_engine.router_ip == expected_host + assert vllm_engine.get_url() == f"http://{expected_host}:8765" + + +@pytest.mark.unit +def test_launch_server_process_brackets_ipv6_health_url(vllm_args, monkeypatch): + process = SimpleNamespace(start=lambda: None, is_alive=lambda: True) + base_urls = [] + subprocess_args = {} + + monkeypatch.setattr(mod, "_build_subprocess_env", lambda _: {}) + monkeypatch.setattr(mod.multiprocessing, "set_start_method", lambda *args, **kwargs: None) + monkeypatch.setattr( + mod.multiprocessing, + "Process", + lambda *, target, args: subprocess_args.update(args[0]) or process, + ) + monkeypatch.setattr(mod, "_wait_server_healthy", lambda base_url, **_: base_urls.append(base_url)) + + launched_process = mod.launch_server_process( + { + "_args": vllm_args, + "_visible_devices": "0", + "host": "[2001:db8::1]", + "port": 8000, + "node_rank": 0, + } + ) + + assert launched_process is process + assert subprocess_args["host"] == "2001:db8::1" + assert base_urls == ["http://[2001:db8::1]:8000"] + + +@pytest.mark.unit +def test_get_base_gpu_id_with_critic_offset(vllm_args): + vllm_args.colocate = False + vllm_args.debug_rollout_only = False + vllm_args.actor_num_gpus_per_node = 4 + vllm_args.actor_num_nodes = 1 + vllm_args.use_critic = True + vllm_args.critic_num_gpus_per_node = 2 + vllm_args.critic_num_nodes = 1 + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 2 + # actor 4 + critic 2 + rank0*2 = 6 + assert mod.get_base_gpu_id(vllm_args, rank=0) == 6 + + +@pytest.mark.unit +def test_resume_memory_occupation_wake_tags_query(vllm_engine, monkeypatch): + seen: list[tuple] = [] + + def fake_post(url, *, params=None, timeout=30, json=None): + seen.append((url, params, timeout, json)) + return _MockResponse(json_data={"ok": True}) + + vllm_args = vllm_engine.args + vllm_args.vllm_enable_sleep_mode = True + monkeypatch.setattr(mod.requests, "post", fake_post) + + vllm_engine.resume_memory_occupation(tags=["weights", "cuda_graph"]) + + assert len(seen) == 1 + assert seen[0][1] == [("tags", "weights")] + + +@pytest.mark.unit +def test_resume_memory_occupation_returns_none_for_unsupported_tags(vllm_engine, monkeypatch): + seen: list[tuple] = [] + + def fake_post(url, *, params=None, timeout=30, json=None): + seen.append((url, params, timeout, json)) + return _MockResponse(json_data={"ok": True}) + + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.resume_memory_occupation(tags=["cuda_graph"]) == {"ok": True} + assert seen[0][1] is None + + +@pytest.mark.unit +def test_release_memory_occupation_flushes_then_posts_sleep(vllm_engine, monkeypatch): + calls: list[str] = [] + + def fake_flush_cache(): + calls.append("flush_cache") + + def fake_post(url, *, params=None, timeout=30, json=None): + calls.append(url) + assert params == {"level": 2} + assert timeout == 30 + assert json is None + return _MockResponse(json_data={"ok": True, "sleep_mode": True}) + + vllm_engine.args.vllm_enable_sleep_mode = False + monkeypatch.setattr(vllm_engine, "flush_cache", fake_flush_cache) + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.release_memory_occupation(level=2) == {"ok": True, "sleep_mode": True} + assert calls == ["flush_cache", "http://127.0.0.1:8765/sleep"] + + +@pytest.mark.unit +def test_resume_memory_occupation_posts_wake_even_when_sleep_disabled(vllm_engine, monkeypatch): + seen: list[tuple] = [] + + def fake_post(url, *, params=None, timeout=30, json=None): + seen.append((url, params, timeout, json)) + return _MockResponse(json_data={"ok": True, "sleep_mode": True}) + + vllm_engine.args.vllm_enable_sleep_mode = False + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.resume_memory_occupation() == {"ok": True, "sleep_mode": True} + assert seen == [("http://127.0.0.1:8765/wake_up", None, 30, None)] + + +@pytest.mark.unit +def test_init_weights_update_group_posts_init_info(vllm_engine, monkeypatch): + calls: list[tuple] = [] + + def fake_post(endpoint: str, payload: dict): + calls.append((endpoint, payload)) + return {"initialized": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_post) + + result = vllm_engine.init_weights_update_group( + "127.0.0.1", + 29500, + rank_offset=1, + world_size=4, + group_name="unused", + backend="nccl", + ) + + assert result == {"initialized": True} + assert len(calls) == 1 + assert calls[0][0] == "init_weight_transfer_engine" + assert calls[0][1]["init_info"]["master_address"] == "127.0.0.1" + assert calls[0][1]["init_info"]["rank_offset"] == 1 + + +@pytest.mark.unit +def test_update_weights_from_disk_posts_collective_rpc(vllm_engine, monkeypatch): + seen: list[tuple] = [] + + def fake_post(url, *, params=None, timeout=30, json=None): + seen.append((url, params, timeout, json)) + return _MockResponse(json_data={"reloaded": True}) + + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8") == {"reloaded": True} + assert seen[0][0] == "http://127.0.0.1:8765/collective_rpc" + assert seen[0][3]["method"] == "reload_weights" + assert seen[1][0] == "http://127.0.0.1:8765/update_weight_version" + assert seen[1][3] == {"new_version": "8"} + assert vllm_engine._weight_version == "8" + + +@pytest.mark.unit +def test_pull_weights_posts_collective_rpc(vllm_engine, monkeypatch): + vllm_engine.args.update_weight_local_checkpoint_dir = "/local/checkpoint" + vllm_engine.args.update_weight_disk_dir = "/shared/checkpoints" + vllm_engine.args.custom_update_weight_pre_read_path = "hooks.refresh" + seen = [] + + def fake_post(url, *, json=None): + seen.append((url, json)) + return _MockResponse(json_data={"success": True, "weight_version": "8"}) + + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.pull_weights(8) == {"success": True, "weight_version": "8"} + assert seen == [ + ( + "http://127.0.0.1:8765/collective_rpc", + { + "method": "pull_weights", + "kwargs": { + "local_checkpoint_dir": "/local/checkpoint", + "source_dir": "/shared/checkpoints", + "target_version": 8, + "pre_read_hook": "hooks.refresh", + }, + }, + ), + ( + "http://127.0.0.1:8765/update_weight_version", + {"new_version": "8"}, + ), + ] + assert vllm_engine._weight_version == "8" + + +@pytest.mark.unit +def test_pull_weights_does_not_advance_version_when_pull_fails(vllm_engine, monkeypatch): + vllm_engine.args.update_weight_local_checkpoint_dir = "/local/checkpoint" + vllm_engine.args.update_weight_disk_dir = "/shared/checkpoints" + vllm_engine.args.custom_update_weight_pre_read_path = None + vllm_engine._weight_version = "old" + + def fake_post(url, *, json=None): + del url, json + return _MockResponse(status_code=500) + + monkeypatch.setattr(mod.requests, "post", fake_post) + + with pytest.raises(requests.exceptions.HTTPError): + vllm_engine.pull_weights(8) + assert vllm_engine._weight_version == "old" + + +@pytest.mark.unit +def test_profile_worker_rank_skips_http(vllm_engine, monkeypatch): + vllm_engine.node_rank = 1 + monkeypatch.setattr(mod.requests, "post", lambda *args, **kwargs: pytest.fail("unexpected HTTP request")) + + assert vllm_engine.start_profile() is None + assert vllm_engine.stop_profile() is None + + +@pytest.mark.unit +def test_update_weights_from_disk_surfaces_http_error(vllm_engine, monkeypatch): + def fake_post(url, *, params=None, timeout=30, json=None): + return _MockResponse(text="boom", status_code=500) + + monkeypatch.setattr(mod.requests, "post", fake_post) + with pytest.raises(requests.exceptions.HTTPError): + vllm_engine.update_weights_from_disk("/tmp/model") + + +@pytest.mark.unit +def test_resolve_parallel_sizes_is_per_engine_not_global(vllm_args): + # The global flag is 1, but THIS engine has 2 GPUs → tp must be 2 (per-engine), not 1. + # A stale global vllm_tp_size must NOT shadow the per-engine value. + vllm_args.rollout_num_gpus_per_engine = 1 + vllm_args.vllm_pipeline_parallel_size = 1 + vllm_args.vllm_tp_size = 1 # stale global; must be ignored now + tp, pp, pcp, dp = mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=2) + assert (tp, pp, pcp, dp) == (2, 1, 1, 1) + + +@pytest.mark.unit +def test_launch_config_heterogeneous_per_group_tp(vllm_args): + vllm_args.num_gpus_per_node = 8 + vllm_args.rollout_num_gpus_per_engine = 1 + vllm_args.vllm_pipeline_parallel_size = 1 + sa, _ = mod._compute_server_args( + vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000, num_gpus_per_engine=2 + ) + assert sa["_tp_size"] == 2 + assert sa["nnodes"] == 1 + + +@pytest.mark.unit +def test_resolve_parallel_sizes_dp_consumes_gpus(vllm_args): + # vLLM DP consumes GPUs (total = tp * pp * pcp * dp), so tp = gpus // (pp * pcp * dp). + # dp=2, pp=1, 4 GPUs/engine → tp=2. + vllm_args.vllm_pipeline_parallel_size = 1 + vllm_args.vllm_data_parallel_size = 2 + vllm_args.vllm_dp_size = 2 + tp, pp, pcp, dp = mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=4) + assert (tp, pp, pcp, dp) == (2, 1, 1, 2) + + +@pytest.mark.unit +def test_resolve_parallel_sizes_dp_and_pp_combined(vllm_args): + # dp=2, pp=2, 8 GPUs/engine → tp = 8 // (2*2) = 2. + vllm_args.vllm_pipeline_parallel_size = 2 + vllm_args.vllm_data_parallel_size = 2 + vllm_args.vllm_dp_size = 2 + tp, pp, pcp, dp = mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=8) + assert (tp, pp, pcp, dp) == (2, 2, 1, 2) + + +@pytest.mark.unit +def test_resolve_parallel_sizes_pcp_consumes_gpus(vllm_args): + vllm_args.vllm_pipeline_parallel_size = 2 + vllm_args.vllm_prefill_context_parallel_size = 2 + vllm_args.vllm_data_parallel_size = 1 + vllm_args.vllm_dp_size = 1 + + tp, pp, pcp, dp = mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=8) + + assert (tp, pp, pcp, dp) == (2, 2, 2, 1) + + +@pytest.mark.unit +@pytest.mark.parametrize("field", ["ep_size", "expert_parallel_size", "moe_dp_size", "moe_data_parallel_size"]) +def test_resolve_parallel_sizes_rejects_pseudo_expert_overrides(vllm_args, field): + with pytest.raises(ValueError, match="does not accept explicit EP/MoE-DP sizes"): + mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=4, overrides={field: 4}) + + +@pytest.mark.unit +@pytest.mark.parametrize("field", ["tp_size", "pp_size", "pcp_size", "dp_size"]) +def test_resolve_parallel_sizes_rejects_non_native_aliases(vllm_args, field): + with pytest.raises(ValueError, match="native field names"): + mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=4, overrides={field: 1}) + + +@pytest.mark.unit +def test_resolve_parallel_sizes_rejects_indivisible_dp(vllm_args): + # gpus_per_engine not divisible by pp*dp must raise (fail fast, not desync the rendezvous). + vllm_args.vllm_pipeline_parallel_size = 1 + vllm_args.vllm_data_parallel_size = 2 + vllm_args.vllm_dp_size = 2 + with pytest.raises(ValueError, match="divisible"): + mod._resolve_parallel_sizes(vllm_args, gpus_per_engine=3) + + +@pytest.mark.unit +def test_make_request_short_circuits_on_headless(vllm_engine, monkeypatch): + # _make_request is the single control-plane POST choke point; on a headless worker + # (node_rank>0) it must no-op to None without issuing any HTTP request. + def _boom(*a, **k): + raise AssertionError("control-plane HTTP must not be called on a headless worker") + + monkeypatch.setattr(mod.requests, "post", _boom) + vllm_engine.node_rank = 1 + assert vllm_engine._make_request("whatever", {}) is None + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tools/convert_hf_to_fp8.py b/tools/convert_hf_to_fp8.py index 5294140d5..52f6df41d 100644 --- a/tools/convert_hf_to_fp8.py +++ b/tools/convert_hf_to_fp8.py @@ -28,6 +28,8 @@ import torch.nn.functional as F from tqdm import tqdm +from vime.utils import accelerator + FP8_INFO = torch.finfo(torch.float8_e4m3fn) FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min @@ -57,7 +59,10 @@ def block_fp8(weight, block_size): block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0] block_max = torch.max(block_max, dim=3, keepdim=True)[0] - scale = block_max.to(torch.float32) / FP8_MAX + # Clamp the block max away from zero, matching channel_fp8 / tensor_fp8: an + # all-zero block (e.g. padding rows or an unused MoE expert) would otherwise + # produce scale == 0 and qweight == 0/0 == NaN in the converted checkpoint. + scale = block_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX qweight = ( (qweight / scale) .clamp(min=FP8_MIN, max=FP8_MAX) @@ -114,11 +119,13 @@ def process_file(input_path, output_path, filename, strategy, block_size, result if not filename.endswith(".safetensors"): return - print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}") + print(f"Processing {filename}, memory usage: {accelerator.memory_allocated()}") weights = {} q_weights = {} - with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f: + with safetensors.safe_open( + os.path.join(input_path, filename), framework="pt", device=accelerator.device_name() + ) as f: for k in f.keys(): weights[k] = f.get_tensor(k) @@ -240,7 +247,7 @@ def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers= json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2) gc.collect() - torch.cuda.empty_cache() + accelerator.empty_cache() if __name__ == "__main__": diff --git a/tools/convert_hf_to_int4_direct.py b/tools/convert_hf_to_int4_direct.py index 613f65595..6b54f46fb 100644 --- a/tools/convert_hf_to_int4_direct.py +++ b/tools/convert_hf_to_int4_direct.py @@ -21,6 +21,8 @@ import torch from tqdm import tqdm +from vime.utils import accelerator + try: import fake_int4_quant_cuda except ImportError: @@ -166,11 +168,13 @@ def add_result(self, filename, q_weights): def process_file(input_path, output_path, filename, group_size, is_symmetric, ignore_rules, result_collector): - print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}") + print(f"Processing {filename}, memory usage: {accelerator.memory_allocated()}") weights = {} q_weights = {} - with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device="cuda") as f: + with safetensors.safe_open( + os.path.join(input_path, filename), framework="pt", device=accelerator.device_name() + ) as f: for k in f.keys(): weights[k] = f.get_tensor(k) @@ -180,15 +184,15 @@ def process_file(input_path, output_path, filename, group_size, is_symmetric, ig ) if is_ignored or not name.endswith(".weight") or weight.dim() < 2: - print(f"Ignoring {name}, memory usage: {torch.cuda.memory_allocated()}") + print(f"Ignoring {name}, memory usage: {accelerator.memory_allocated()}") q_weights[name] = weight continue - print(f"Packing {name}, memory usage: {torch.cuda.memory_allocated()}") + print(f"Packing {name}, memory usage: {accelerator.memory_allocated()}") qw, s, zp = pack_layer(weight, group_size, is_symmetric) qweight_name = name.replace(".weight", ".weight_packed") scale_name = name.replace(".weight", ".weight_scale") - weight_shape = torch.tensor(weight.shape, dtype=torch.int32, device="cuda") + weight_shape = torch.tensor(weight.shape, dtype=torch.int32, device=accelerator.device()) weight_shape_name = name.replace(".weight", ".weight_shape") if zp is not None: zp_name = name.replace(".weight", ".weight_zero_point") @@ -273,7 +277,7 @@ def convert_int4(input_path, output_path, group_size, is_symmetric, ignore_rules json.dump(index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2) gc.collect() - torch.cuda.empty_cache() + accelerator.empty_cache() return output_path diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 117969fa4..45c7360ac 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -1,24 +1,26 @@ +import argparse import gc import os import shutil import torch import torch.distributed as dist -from vime.utils.common import is_npu +from vime.platforms import current_platform + +if current_platform().is_npu: + import vime.backends.megatron_utils # noqa: F401 -if is_npu(): - import megatron_adaptor # noqa: F401 from megatron.core.enums import ModelType from megatron.training.arguments import parse_args, validate_args from megatron.training.checkpointing import get_checkpoint_name, get_checkpoint_tracker_filename, save_checkpoint from megatron.training.training import get_model -import vime_plugins.mbridge # noqa: F401 -from mbridge import AutoBridge from vime.backends.megatron_utils.arguments import set_default_megatron_args +from vime.backends.megatron_utils.hf_to_megatron import load_hf_weights from vime.backends.megatron_utils.initialize import init from vime.backends.megatron_utils.model_provider import get_model_provider_func -from vime.utils.logging_utils import configure_logger +from vime.observability.logging_utils import configure_logger +from vime.utils import accelerator from vime.utils.memory_utils import print_memory @@ -26,11 +28,16 @@ def add_convertion_args(parser): """Add conversion arguments to the parser""" parser.add_argument("--hf-checkpoint", type=str, required=True, help="HuggingFace model path") parser.add_argument( - "--megatron-to-hf-mode", - choices=["raw", "bridge"], - default="raw", - help="The method to convert megatron weights to hugging face weights for vLLM.", + "--custom-model-provider-path", + type=str, + default=None, + help="Path to a custom model provider function.", ) + parser.add_argument("--allgather-cp", action="store_true", default=False) + try: + parser.add_argument("--use-gated-attention", action="store_true", default=False) + except argparse.ArgumentError: + pass try: parser.add_argument("--padded-vocab-size", type=int, default=None) except Exception: @@ -82,6 +89,14 @@ def ceildiv(a, b): def main(): + if torch.version.hip: + import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module + + from vime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync + + filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync + print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility") + configure_logger() # Initialize distributed environment @@ -89,25 +104,18 @@ def main(): local_rank = int(os.getenv("LOCAL_RANK") or os.getenv("SLURM_LOCALID") or 0) global_rank = int(os.getenv("RANK") or os.getenv("SLURM_PROCID") or 0) - torch.cuda.set_device(local_rank) + accelerator.set_device(local_rank) os.environ.setdefault("WORLD_SIZE", str(world_size)) os.environ.setdefault("RANK", str(global_rank)) os.environ.setdefault("LOCAL_RANK", str(local_rank)) os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "12355") - if is_npu(): - dist.init_process_group( - backend="hccl", - world_size=world_size, - rank=global_rank, - ) - else: - dist.init_process_group( - backend="nccl", - world_size=world_size, - rank=global_rank, - device_id=torch.device(f"cuda:{local_rank}"), - ) + dist.init_process_group( + backend=accelerator.process_group_backend(), + world_size=world_size, + rank=global_rank, + device_id=accelerator.distributed_device_id(local_rank), + ) args = get_args() init(args) @@ -115,17 +123,16 @@ def main(): # Load model hf_model_path = args.hf_checkpoint - bridge = AutoBridge.from_pretrained(hf_model_path, trust_remote_code=True) - bridge.load_weights(model, hf_model_path, memory_efficient=True) + load_hf_weights(args, model, hf_model_path) print(f"Model loaded: {hf_model_path}") if args.use_cpu_initialization: model[0] = model[0].cpu() print_memory("after loading model") - torch.cuda.synchronize() + accelerator.synchronize() gc.collect() - torch.cuda.empty_cache() + accelerator.empty_cache() save_checkpoint(1, model, None, None, 0) diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py index 62308239e..5f3aa221e 100644 --- a/tools/convert_to_hf.py +++ b/tools/convert_to_hf.py @@ -5,6 +5,7 @@ import vime.backends.megatron_utils as megatron_utils from vime.backends.megatron_utils import update_weight_utils +from vime.utils import accelerator from vime.utils.arguments import parse_args @@ -57,7 +58,7 @@ def main(args): param = param_ break else: - param = torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()) + param = torch.empty(info.shape, dtype=info.dtype, device=accelerator.current_device()) if pp_size > 1: if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()): diff --git a/tools/convert_torch_dist_to_hf_bridge.py b/tools/convert_torch_dist_to_hf_bridge.py deleted file mode 100644 index 7c67bf018..000000000 --- a/tools/convert_torch_dist_to_hf_bridge.py +++ /dev/null @@ -1,65 +0,0 @@ -import argparse -import os - -import megatron.bridge.training.model_load_save as _model_load_save_module -from megatron.bridge import AutoBridge - -from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config - - -# Here we need to patch Megatron Bridge's `load_model_config`, since the checkpoint is saved -# by Megatron and lack of provider information. -_provider_override = {} -_original_load_model_config = _model_load_save_module.load_model_config - - -def _patched_load_model_config(checkpoint_path): - model_cfg, mlm_args = _original_load_model_config(checkpoint_path) - provider = _provider_override.get("provider") - if provider is not None: - from megatron.bridge.models.model_provider import ModelProviderMixin - - if not isinstance(model_cfg, ModelProviderMixin): - print(f"[convert] Overriding MLM TransformerConfig with Bridge provider: " f"{type(provider).__name__}") - return provider, mlm_args - return model_cfg, mlm_args - - -_model_load_save_module.load_model_config = _patched_load_model_config - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Convert torch distributed checkpoint to HuggingFace format using Megatron Bridge" - ) - parser.add_argument( - "--input-dir", type=str, required=True, help="Path to the torch distributed checkpoint directory" - ) - parser.add_argument("--output-dir", type=str, required=True, help="Path to save the HuggingFace checkpoint") - parser.add_argument( - "--origin-hf-dir", - type=str, - required=True, - help="Path to the original HuggingFace model directory (for config)", - ) - parser.add_argument( - "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists." - ) - args = parser.parse_args() - - if os.path.exists(args.output_dir) and not args.force: - raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.") - - print(f"Loading config from {args.origin_hf_dir}") - bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.origin_hf_dir, trust_remote_code=True)) - - # Use Bridge's provider so the correct model class is created (e.g., Qwen3VLModel - # instead of GPTModel). This is needed because MLM checkpoints lack run_config.yaml. - provider = bridge.to_megatron_provider(load_weights=False) - _provider_override["provider"] = provider - print(f"[convert] Using Bridge provider: {type(provider).__name__}") - - print(f"Exporting checkpoint from {args.input_dir} to {args.output_dir}") - bridge.export_ckpt(args.input_dir, args.output_dir) - - print("Done!") diff --git a/tools/fp8_cast_bf16.py b/tools/fp8_cast_bf16.py index c227c300f..3862ed5dd 100644 --- a/tools/fp8_cast_bf16.py +++ b/tools/fp8_cast_bf16.py @@ -10,6 +10,8 @@ from safetensors.torch import load_file, save_file from tqdm import tqdm +from vime.utils import accelerator + @triton.jit def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): @@ -60,7 +62,7 @@ def get_tensor(tensor_name): file_name = weight_map[tensor_name] if file_name not in loaded_files: file_path = os.path.join(fp8_path, file_name) - loaded_files[file_name] = load_file(file_path, device="cuda") + loaded_files[file_name] = load_file(file_path, device=accelerator.device_name()) return loaded_files[file_name][tensor_name] safetensor_files = list(glob(os.path.join(fp8_path, "*.safetensors"))) @@ -68,7 +70,7 @@ def get_tensor(tensor_name): for safetensor_file in tqdm(safetensor_files): print(f"Handling file: {safetensor_file}") file_name = os.path.basename(safetensor_file) - current_state_dict = load_file(safetensor_file, device="cuda") + current_state_dict = load_file(safetensor_file, device=accelerator.device_name()) loaded_files[file_name] = current_state_dict new_state_dict = {} @@ -95,7 +97,7 @@ def get_tensor(tensor_name): if len(loaded_files) > 2: oldest_file = next(iter(loaded_files)) del loaded_files[oldest_file] - torch.cuda.empty_cache() + accelerator.empty_cache() # Update model index new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json") diff --git a/tools/preprocess_gpt_oss.py b/tools/preprocess_gpt_oss.py deleted file mode 100644 index 9b12c342e..000000000 --- a/tools/preprocess_gpt_oss.py +++ /dev/null @@ -1,266 +0,0 @@ -"""Preprocess GPT-OSS model: dequantize MXFP4 experts and unfuse into per-expert format. - -This converts the GPT-OSS HF checkpoint from: - - MXFP4 quantized fused expert weights (gate_up_proj_blocks/scales, down_proj_blocks/scales) -To: - - BF16 per-expert weights (experts.{e}.gate_proj.weight, experts.{e}.up_proj.weight, etc.) - -Usage: - python tools/preprocess_gpt_oss.py \ - --input /path/to/gpt-oss-20b \ - --output /path/to/gpt-oss-20b-bf16 -""" - -import argparse -import json -import math -import os -import shutil -from collections import OrderedDict - -import torch -from safetensors import safe_open -from safetensors.torch import save_file - - -def dequantize_mxfp4( - blocks: torch.Tensor, - scales: torch.Tensor, - dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - """Dequantize MXFP4 weights to BF16. Adapted from megatron.bridge.""" - FP4_VALUES = [ - +0.0, - +0.5, - +1.0, - +1.5, - +2.0, - +3.0, - +4.0, - +6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ] - scales = scales.to(torch.int32) - 127 - lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device) - - *prefix_shape, G, B = blocks.shape - rows_total = math.prod(prefix_shape) * G - - blocks_flat = blocks.reshape(rows_total, B) - scales_flat = scales.reshape(rows_total, 1) - - out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device) - - rows_per_chunk = 32768 * 1024 - for r0 in range(0, rows_total, rows_per_chunk): - r1 = min(r0 + rows_per_chunk, rows_total) - blk = blocks_flat[r0:r1] - exp = scales_flat[r0:r1] - - idx_lo = (blk & 0x0F).to(torch.long) - idx_hi = (blk >> 4).to(torch.long) - - sub = out[r0:r1] - sub[:, 0::2] = lut[idx_lo] - sub[:, 1::2] = lut[idx_hi] - torch.ldexp(sub, exp, out=sub) - - return out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2) - - -def preprocess_gpt_oss(input_dir: str, output_dir: str): - os.makedirs(output_dir, exist_ok=True) - - # Load config - with open(os.path.join(input_dir, "config.json")) as f: - config = json.load(f) - - num_experts = config["num_local_experts"] - intermediate_size = config["intermediate_size"] - - # Remove quantization config and ensure torch_dtype is bfloat16 - new_config = {k: v for k, v in config.items() if k != "quantization_config"} - new_config["torch_dtype"] = "bfloat16" - with open(os.path.join(output_dir, "config.json"), "w") as f: - json.dump(new_config, f, indent=2) - - # Copy non-weight files - for fname in os.listdir(input_dir): - if fname in ("config.json", "model.safetensors.index.json"): - continue - if fname.endswith(".safetensors"): - continue - src = os.path.join(input_dir, fname) - dst = os.path.join(output_dir, fname) - if os.path.isfile(src) and not os.path.exists(dst): - shutil.copy2(src, dst) - - # Process safetensors: collect all weight names - index_path = os.path.join(input_dir, "model.safetensors.index.json") - if os.path.exists(index_path): - with open(index_path) as f: - index = json.load(f) - weight_map = index["weight_map"] - else: - # Single file - weight_map = None - - # Group weights by safetensors file - if weight_map: - files = set(weight_map.values()) - else: - files = [f for f in os.listdir(input_dir) if f.endswith(".safetensors")] - - all_output_tensors = OrderedDict() - new_weight_map = {} - - for sf_file in sorted(files): - sf_path = os.path.join(input_dir, sf_file) - print(f"Processing {sf_file}...") - with safe_open(sf_path, framework="pt", device="cpu") as f: - keys = list(f.keys()) - for key in keys: - tensor = f.get_tensor(key) - - # Check if this is a quantized expert weight - if key.endswith("_blocks"): - base_name = key[: -len("_blocks")] - scales_key = base_name + "_scales" - # Load scales from same or different file - try: - scales = f.get_tensor(scales_key) - except Exception: - # Try loading from other files - scales = _load_tensor_from_files(input_dir, files, scales_key) - - print(f" Dequantizing {base_name}...") - dequantized = dequantize_mxfp4(tensor, scales) - _unfuse_experts( - base_name, dequantized, num_experts, intermediate_size, all_output_tensors, new_weight_map - ) - - elif key.endswith("_scales"): - continue # handled with _blocks - - elif ".mlp.experts.gate_up_proj_bias" in key: - # Unfuse bias: [E, 2*intermediate] -> per-expert gate_bias + up_bias - # GPT-OSS uses interleaved format: even=gate, odd=up - layer_prefix = key.rsplit(".mlp.experts.gate_up_proj_bias", 1)[0] - for e in range(num_experts): - gate_bias = tensor[e, 0::2] - up_bias = tensor[e, 1::2] - gname = f"{layer_prefix}.mlp.experts.{e}.gate_proj.bias" - uname = f"{layer_prefix}.mlp.experts.{e}.up_proj.bias" - all_output_tensors[gname] = gate_bias.contiguous() - all_output_tensors[uname] = up_bias.contiguous() - new_weight_map[gname] = "model.safetensors" - new_weight_map[uname] = "model.safetensors" - - elif ".mlp.experts.down_proj_bias" in key: - # Unfuse bias: [E, hidden] -> per-expert - layer_prefix = key.rsplit(".mlp.experts.down_proj_bias", 1)[0] - for e in range(num_experts): - dname = f"{layer_prefix}.mlp.experts.{e}.down_proj.bias" - all_output_tensors[dname] = tensor[e].contiguous() - new_weight_map[dname] = "model.safetensors" - - else: - all_output_tensors[key] = tensor - new_weight_map[key] = "model.safetensors" - - # Save output - print(f"Saving {len(all_output_tensors)} tensors...") - - # Split into chunks of ~5GB each - chunk_size = 5 * 1024 * 1024 * 1024 # 5GB - chunks = [] - current_chunk = OrderedDict() - current_size = 0 - - for name, tensor in all_output_tensors.items(): - tensor_size = tensor.numel() * tensor.element_size() - if current_size + tensor_size > chunk_size and current_chunk: - chunks.append(current_chunk) - current_chunk = OrderedDict() - current_size = 0 - current_chunk[name] = tensor - current_size += tensor_size - - if current_chunk: - chunks.append(current_chunk) - - final_weight_map = {} - if len(chunks) == 1: - out_file = "model.safetensors" - save_file(chunks[0], os.path.join(output_dir, out_file)) - for k in chunks[0]: - final_weight_map[k] = out_file - else: - total = len(chunks) - for i, chunk in enumerate(chunks): - out_file = f"model-{i+1:05d}-of-{total:05d}.safetensors" - save_file(chunk, os.path.join(output_dir, out_file)) - for k in chunk: - final_weight_map[k] = out_file - - # Save index - total_size = sum(t.numel() * t.element_size() for t in all_output_tensors.values()) - index_data = { - "metadata": {"total_size": total_size}, - "weight_map": final_weight_map, - } - with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f: - json.dump(index_data, f, indent=2) - - print(f"Done! Output saved to {output_dir}") - - -def _unfuse_experts(base_name, dequantized, num_experts, intermediate_size, output_tensors, weight_map): - """Unfuse 3D expert tensor into per-expert format.""" - layer_prefix = base_name.rsplit(".mlp.experts.", 1)[0] - weight_type = base_name.rsplit(".mlp.experts.", 1)[1] - - if "gate_up_proj" in weight_type: - # [E, 2*intermediate, hidden] -> per-expert gate + up - for e in range(num_experts): - expert_weight = dequantized[e] # [2*intermediate, hidden] - # GPT-OSS uses interleaved format: [g0,u0,g1,u1,...] (even=gate, odd=up) - gate = expert_weight[0::2] # [intermediate, hidden] - up = expert_weight[1::2] # [intermediate, hidden] - gname = f"{layer_prefix}.mlp.experts.{e}.gate_proj.weight" - uname = f"{layer_prefix}.mlp.experts.{e}.up_proj.weight" - output_tensors[gname] = gate.contiguous() - output_tensors[uname] = up.contiguous() - weight_map[gname] = "model.safetensors" - weight_map[uname] = "model.safetensors" - elif "down_proj" in weight_type: - # [E, hidden, intermediate] -> per-expert - for e in range(num_experts): - dname = f"{layer_prefix}.mlp.experts.{e}.down_proj.weight" - output_tensors[dname] = dequantized[e].contiguous() - weight_map[dname] = "model.safetensors" - - -def _load_tensor_from_files(input_dir, files, key): - """Load a tensor by searching across safetensors files.""" - for sf_file in files: - sf_path = os.path.join(input_dir, sf_file) - with safe_open(sf_path, framework="pt", device="cpu") as f: - if key in f.keys(): - return f.get_tensor(key) - raise KeyError(f"Tensor {key} not found in any safetensors file") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Preprocess GPT-OSS model") - parser.add_argument("--input", required=True, help="Input HF model directory") - parser.add_argument("--output", required=True, help="Output directory for BF16 model") - args = parser.parse_args() - preprocess_gpt_oss(args.input, args.output) diff --git a/tools/profile_rollout.py b/tools/profile_rollout.py index c3802d2eb..0709e554a 100644 --- a/tools/profile_rollout.py +++ b/tools/profile_rollout.py @@ -14,18 +14,10 @@ def get_workers(router_url): return [] -def start_profile(worker_url, args): - payload = { - "output_dir": args.output_dir, - "num_steps": args.num_steps, - "activities": args.activities, - "profile_by_stage": args.profile_by_stage, - "with_stack": args.with_stack, - "record_shapes": args.record_shapes, - } +def start_profile(worker_url): try: - print(f"Starting profile on {worker_url} for {args.num_steps} steps...") - response = requests.post(f"{worker_url}/start_profile", json=payload) + print(f"Starting profile on {worker_url}...") + response = requests.post(f"{worker_url}/start_profile", json={}) response.raise_for_status() print(f"Successfully started profile on {worker_url}") except Exception as e: @@ -46,12 +38,6 @@ def main(): parser = argparse.ArgumentParser(description="Automate vLLM profiling across all workers via router.") parser.add_argument("--router-url", type=str, required=True, help="Router URL (e.g., http://127.0.0.1:3000)") parser.add_argument("--action", type=str, choices=["start", "stop"], default="start", help="Action to perform") - parser.add_argument("--output-dir", type=str, default="/tmp/vllm_profile", help="Output directory for traces") - parser.add_argument("--num-steps", type=int, default=3, help="Number of steps to profile (default: 3)") - parser.add_argument("--activities", type=str, nargs="+", default=["GPU"], help="Activities to profile (CPU, GPU)") - parser.add_argument("--profile-by-stage", action="store_true", help="Profile by stage (prefill/decode)") - parser.add_argument("--with-stack", action="store_true", help="Record call stack") - parser.add_argument("--record-shapes", action="store_true", help="Record tensor shapes") args = parser.parse_args() @@ -68,7 +54,7 @@ def main(): continue if args.action == "start": - start_profile(worker_url, args) + start_profile(worker_url) else: stop_profile(worker_url) diff --git a/tools/trace_timeline_viewer.py b/tools/trace_timeline_viewer.py index f08683f2e..274c784bb 100644 --- a/tools/trace_timeline_viewer.py +++ b/tools/trace_timeline_viewer.py @@ -271,7 +271,7 @@ def _build_items_from_trace(sample: dict[str, Any], sample_idx: int) -> dict[str "display_end_ts": None, "attempt": event["attempt"], "span_id": event["span_id"], - "parent_span_id": event.get("parent_span_id"), + "parent_span_id": event.get("parent_span_id") or event.get("inferred_parent_span_id"), "start_attrs": event.get("attrs") or {}, "end_attrs": {}, } @@ -309,7 +309,7 @@ def _build_items_from_trace(sample: dict[str, Any], sample_idx: int) -> dict[str "ts": event["ts"], "attempt": event["attempt"], "span_id": None, - "parent_span_id": event.get("inferred_parent_span_id"), + "parent_span_id": event.get("parent_span_id") or event.get("inferred_parent_span_id"), "attrs": event.get("attrs") or {}, } ) @@ -376,12 +376,12 @@ def nearest_closed_ancestor_end(span: dict[str, Any]) -> float | None: for event in point_events: parent_span_id = event.get("parent_span_id") - event["depth"] = span_depths.get(parent_span_id or "", 0) + event["depth"] = span_depths[parent_span_id] + 1 if parent_span_id in span_depths else 0 event["lane"] = event["depth"] for item in orphan_ends: parent_span_id = item.get("parent_span_id") - item["depth"] = span_depths.get(parent_span_id or "", 0) + item["depth"] = span_depths[parent_span_id] + 1 if parent_span_id in span_depths else 0 item["lane"] = item["depth"] def parent_span_name(parent_span_id: str | None) -> str | None: @@ -455,8 +455,8 @@ def parent_span_name(parent_span_id: str | None) -> str | None: "P", [ "pd_prefill_bootstrap_queue_duration", - "pd_bootstrap_duration", - "pd_alloc_waiting_duration", + "pd_prefill_bootstrap_duration", + "pd_prefill_alloc_wait_duration", "pd_prefill_forward_duration", "pd_prefill_transfer_queue_duration", ], @@ -466,6 +466,8 @@ def parent_span_name(parent_span_id: str | None) -> str | None: "D", [ "pd_decode_prealloc_duration", + "pd_decode_bootstrap_duration", + "pd_decode_alloc_wait_duration", "pd_decode_transfer_duration", "pd_decode_forward_duration", ], @@ -475,6 +477,8 @@ def parent_span_name(parent_span_id: str | None) -> str | None: for span in all_spans: if span["state"] != "closed_span" or span.get("end_ts") is None: continue + if str(span.get("name") or "").startswith("vllm_pd_"): + continue end_attrs = span.get("end_attrs") or {} for role, suffix, keys in pd_lane_specs: role_attrs = { @@ -1243,13 +1247,15 @@ def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]: ['prefill_forward', '[P] prefill fwd'], ['prefill_bootstrap_queue', '[P] bootstrap queue'], ['prefill_transfer_queue', '[P] transfer'], - ['bootstrap', '[P] bootstrap'], - ['alloc_waiting', '[P] alloc wait'], + ['prefill_bootstrap', '[P] bootstrap'], + ['prefill_alloc_wait', '[P] alloc wait'], ['decode_forward', '[D] decode fwd'], ['decode_transfer', '[D] kv transfer'], + ['decode_bootstrap', '[D] bootstrap'], + ['decode_alloc_wait', '[D] alloc wait'], ['decode_prealloc', '[D] prealloc'], ]; - html += 'PD (collapsed: P over D; expanded: separate P/D lanes):'; + html += 'PD phases:'; for (const [phase, label] of pdLegend) { html += `` + `` + @@ -1420,11 +1426,13 @@ def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]: // PD phase color palette (warm earth tones matching the viewer aesthetic) const PD_COLORS = { prefill_bootstrap_queue: 'rgba(180, 160, 120, A)', // muted sand - bootstrap: 'rgba(140, 120, 90, A)', // dark sand - alloc_waiting: 'rgba(160, 140, 110, A)', // warm grey + prefill_bootstrap: 'rgba(140, 120, 90, A)', // dark sand + prefill_alloc_wait: 'rgba(160, 140, 110, A)', // warm grey prefill_forward: 'rgba(70, 140, 180, A)', // steel blue (prefill) prefill_transfer_queue: 'rgba(200, 160, 60, A)', // amber (transfer) decode_prealloc: 'rgba(160, 140, 110, A)', // warm grey + decode_bootstrap: 'rgba(140, 120, 90, A)', // dark sand + decode_alloc_wait: 'rgba(160, 140, 110, A)', // warm grey decode_transfer: 'rgba(200, 160, 60, A)', // amber (transfer) decode_forward: 'rgba(80, 170, 100, A)', // green (decode) }; @@ -1452,12 +1460,14 @@ def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]: }; // P-side phases (sequential order) push('prefill_bootstrap_queue', 'prefill_bootstrap_queue_duration'); - push('bootstrap', 'bootstrap_duration'); - push('alloc_waiting', 'alloc_waiting_duration'); + push('prefill_bootstrap', 'prefill_bootstrap_duration'); + push('prefill_alloc_wait', 'prefill_alloc_wait_duration'); push('prefill_forward', 'prefill_forward_duration'); push('prefill_transfer_queue', 'prefill_transfer_queue_duration'); // D-side phases (sequential order) push('decode_prealloc', 'decode_prealloc_duration'); + push('decode_bootstrap', 'decode_bootstrap_duration'); + push('decode_alloc_wait','decode_alloc_wait_duration'); push('decode_transfer', 'decode_transfer_duration'); push('decode_forward', 'decode_forward_duration'); return segs.length > 0 ? segs : null; @@ -1473,8 +1483,7 @@ def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]: totalDuration: 0, }; } - const pPhases = phases.filter((p) => - p.phase.startsWith('prefill_') || p.phase === 'bootstrap' || p.phase === 'alloc_waiting'); + const pPhases = phases.filter((p) => p.phase.startsWith('prefill_')); const dPhases = phases.filter((p) => p.phase.startsWith('decode_')); const pTotal = pPhases.reduce((sum, p) => sum + p.duration, 0); const dTotal = dPhases.reduce((sum, p) => sum + p.duration, 0); @@ -1880,7 +1889,7 @@ def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]: const phases = pdPhases(attrs); if (phases) { const totalDur = phases.reduce((s, p) => s + p.duration, 0); - const pPhases = phases.filter(p => p.phase.startsWith('prefill_') || p.phase === 'bootstrap' || p.phase === 'alloc_waiting'); + const pPhases = phases.filter(p => p.phase.startsWith('prefill_')); const dPhases = phases.filter(p => p.phase.startsWith('decode_')); const fmtPhase = (p) => { const ms = (p.duration * 1000).toFixed(1); diff --git a/train.py b/train.py index 949ef1609..9b17f2ff4 100644 --- a/train.py +++ b/train.py @@ -1,17 +1,20 @@ import ray +from vime.platforms import current_platform + +if current_platform().is_npu: + import vime.backends.megatron_utils # noqa: F401 + +from vime.observability.logging_utils import configure_logger, finish_tracking, init_tracking from vime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from vime.utils.arguments import parse_args -from vime.utils.common import is_npu -from vime.utils.logging_utils import configure_logger, finish_tracking, init_tracking, update_tracking_open_metrics from vime.utils.misc import should_run_periodic_action -if is_npu(): - import megatron_adaptor # noqa: F401 - def train(args): configure_logger() + release_train = args.release_train + # allocate the GPUs pgs = create_placement_groups(args) init_tracking(args) @@ -20,14 +23,9 @@ def train(args): # need to initialize rollout manager first to calculate num_rollout rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) - # Update primary W&B with vLLM metrics endpoint now that servers are up. - router_addr = ray.get(rollout_manager.get_metrics_router_addr.remote()) - update_tracking_open_metrics(args, router_addr) - - # create the actor and critic models actor_model, critic_model = create_training_models(args, pgs, rollout_manager) - if args.offload_rollout: + if args.offload_rollout and not release_train: ray.get(rollout_manager.onload_weights.remote()) # Always push actor weights to rollout once weights are loaded. @@ -52,21 +50,6 @@ def offload_train(actor_trains_this_step): else: critic_model.clear_memory() - def save(rollout_id): - actor_trains_this_step = (not args.use_critic) or rollout_id >= args.num_critic_only_steps - if actor_trains_this_step: - actor_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) - if args.use_critic: - critic_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) - if args.rollout_global_dataset: - ray.get(rollout_manager.save.remote(rollout_id)) - # train loop. for rollout_id in range(args.start_rollout_id, args.num_rollout): if args.eval_interval is not None and rollout_id == 0 and not args.skip_eval_before_train: @@ -77,22 +60,32 @@ def save(rollout_id): if args.offload_rollout: ray.get(rollout_manager.offload.remote()) - actor_trains_this_step = (not args.use_critic) or rollout_id >= args.num_critic_only_steps + if release_train: + actor_model.create() + actor_trains = (not args.use_critic) or rollout_id >= args.num_critic_only_steps if args.use_critic: value_refs = critic_model.async_train(rollout_id, rollout_data_ref) - if actor_trains_this_step: + if actor_trains: ray.get(actor_model.async_train(rollout_id, rollout_data_ref, external_data=value_refs)) else: ray.get(value_refs) else: ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) - if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): - save(rollout_id) - - offload_train(actor_trains_this_step) - if args.offload_rollout: + if release_train or should_run_periodic_action( + rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout + ): + force_sync = release_train or rollout_id == args.num_rollout - 1 + if actor_trains: + actor_model.save_model(rollout_id, force_sync=force_sync) + if args.use_critic: + critic_model.save_model(rollout_id, force_sync=force_sync) + if args.rollout_global_dataset: + ray.get(rollout_manager.save.remote(rollout_id)) + + offload_train(actor_trains) + if args.offload_rollout and not release_train: ray.get(rollout_manager.onload_weights.remote()) actor_model.update_weights() diff --git a/train_async.py b/train_async.py index 93aad67b3..5d58ecfb6 100644 --- a/train_async.py +++ b/train_async.py @@ -1,19 +1,21 @@ import ray +from vime.platforms import current_platform + +if current_platform().is_npu: + import vime.backends.megatron_utils # noqa: F401 + +from vime.observability.logging_utils import configure_logger, finish_tracking, init_tracking from vime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from vime.utils.arguments import parse_args -from vime.utils.common import is_npu -from vime.utils.logging_utils import configure_logger, finish_tracking, init_tracking, update_tracking_open_metrics from vime.utils.misc import should_run_periodic_action -if is_npu(): - import megatron_adaptor # noqa: F401 - # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). def train(args): assert not args.colocate, "Colocation is not supported for async training." configure_logger() + release_train = args.release_train # allocate the GPUs pgs = create_placement_groups(args) init_tracking(args) @@ -22,10 +24,6 @@ def train(args): # need to initialize rollout manager first to calculate num_rollout rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) - # Update primary W&B with vLLM metrics endpoint now that servers are up. - router_addr = ray.get(rollout_manager.get_metrics_router_addr.remote()) - update_tracking_open_metrics(args, router_addr) - # create the actor and critic models actor_model, critic_model = create_training_models(args, pgs, rollout_manager) @@ -46,31 +44,31 @@ def train(args): if rollout_id + 1 < args.num_rollout: rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1) + if release_train: + actor_model.create() + + actor_trains = (not args.use_critic) or rollout_id >= args.num_critic_only_steps if args.use_critic: - actor_trains_this_step = rollout_id >= args.num_critic_only_steps value_refs = critic_model.async_train(rollout_id, rollout_data_curr_ref) - if actor_trains_this_step: + if actor_trains: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref, external_data=value_refs)) else: ray.get(value_refs) else: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) - if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): - if (not args.use_critic) or rollout_id >= args.num_critic_only_steps: - actor_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) + if release_train or should_run_periodic_action( + rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout + ): + force_sync = release_train or rollout_id == args.num_rollout - 1 + if actor_trains: + actor_model.save_model(rollout_id, force_sync=force_sync) if args.use_critic: - critic_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) + critic_model.save_model(rollout_id, force_sync=force_sync) if args.rollout_global_dataset: ray.get(rollout_manager.save.remote(rollout_id)) - if (rollout_id + 1) % args.update_weights_interval == 0: + if release_train or (rollout_id + 1) % args.update_weights_interval == 0: # sync generate before update weights to prevent update weight in the middle of generation rollout_data_curr_ref = ray.get(x) if (x := rollout_data_next_future) is not None else None rollout_data_next_future = None diff --git a/vime/agent/adapters/anthropic.py b/vime/agent/adapters/anthropic.py index 2d3c23492..8bbdbf743 100644 --- a/vime/agent/adapters/anthropic.py +++ b/vime/agent/adapters/anthropic.py @@ -1,20 +1,22 @@ """Anthropic Messages adapter for agent rollouts. -The adapter exposes ``/v1/messages`` and ``/v1/messages/count_tokens``. It -renders each Anthropic message history with the served model's chat template, -calls vLLM's ``/inference/v1/generate`` with ``token_ids``, and -records the exact sampled token ids/logprobs as ``TurnRecord`` objects. New -code should use ``AnthropicAdapter`` and call ``finish_session()`` at trajectory -end to drain trainable ``TokenSegment`` objects. - -It also handles Claude Code sub-agent and compaction patterns by splitting one -session into ``subagent``, ``wipe``, and ``final`` segments. +Exposes /v1/messages and /v1/messages/count_tokens. Each Anthropic message +history is rendered with the served model's chat template, sent to vllm +``/inference/v1/generate`` as ``token_ids``, and fed into a shared +TrajectoryManager keyed by session id. finish_session(sid) drains a session's +trajectory into a list of Sample. + +The per-sid tree inside TrajectoryManager handles sub-agent and compaction +patterns automatically: any divergence in the prompt prefix forks into a new +leaf, so we do not track explicit chains here. + +This module mirrors vime.agent.adapters.openai; the section layout (adapter +class -> translation -> reply building -> request framing) is shared between +them. See BaseAdapter for the hooks to fill. """ from __future__ import annotations -import asyncio -import dataclasses import json import logging import secrets @@ -22,184 +24,82 @@ from aiohttp import web -from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY -from vime.agent.adapters.common import AdapterChain as Chain from vime.agent.adapters.common import ( BaseAdapter, - call_vllm_generate, - ok_response, - render_token_ids, - request_session_id, + Reply, + flatten_content, + manager_finish_reason, + sid_from_bearer, + tool_call_dict, ) -from vime.agent.adapters.common import stable_hash as _hash -from vime.agent.parsing import parse_model_output -from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments +from vime.agent.parsing import ParsedModelOutput logger = logging.getLogger(__name__) -# Tool names claude-code uses to dispatch a sub-agent. -_SUBAGENT_TOOLS = {"Task", "Agent"} - - -@dataclasses.dataclass -class Session: - main: Chain = dataclasses.field(default_factory=Chain) - active_sub: Chain | None = None # at most one sub-agent at a time - pending_dispatch_id: str = "" # tool_use_id we're waiting to close - sampling_defaults: dict = dataclasses.field(default_factory=dict) - max_context_tokens: int = 0 - lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - segments: list[TurnSegment] = dataclasses.field(default_factory=list) # frozen output - - class AnthropicAdapter(BaseAdapter): - """Anthropic Messages-compatible HTTP adapter with session lifecycle helpers.""" - - session_cls = Session - - def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: - super().__init__( - tokenizer=tokenizer, - vllm_url=vllm_url, - tool_parser=tool_parser, - reasoning_parser=reasoning_parser, + """Anthropic Messages-compatible HTTP adapter: wire translation and reply + framing only; the turn machinery is inherited from BaseAdapter.""" + + logger = logger + log_prefix = "anthropic_adapter" + max_token_keys = ("max_tokens",) + stop_keys = ("stop_sequences",) + + def _register_routes(self, app: web.Application) -> None: + app.router.add_post("/v1/messages", self._run_turn) + app.router.add_post("/v1/messages/count_tokens", _count_tokens) + + def _session_id(self, request: web.Request, body: dict) -> str: + return _request_session_id(request) + + def _preprocess_body(self, body: dict) -> None: + _fold_mid_list_system_into_user(body) + + def _translate(self, body: dict) -> tuple[list[dict], list[dict] | None]: + translated = _translate_messages(body.get("messages") or [], body.get("system")) + tools_schema = _tools_to_chat_tools(body.get("tools")) + return translated, tools_schema + + def _build_reply(self, parsed, raw_finish, translated, tools_schema) -> Reply: + blocks, stop_reason, manager_message = _build_reply_parts(parsed, raw_finish) + return Reply( + manager_message=manager_message, + finish_reason=manager_finish_reason(parsed.tool_uses, raw_finish), + wire=(blocks, stop_reason), ) - self.app.router.add_post("/v1/messages", _handle_request) - self.app.router.add_post("/v1/messages/count_tokens", _count_tokens) - self.app.router.add_get("/healthz", _ok) - self.app.router.add_get("/v1/models", _ok) - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - await self.shutdown_session(sid, wait_timeout=wait_timeout) - s = self.store.pop(sid, None) - if s is None: - return [] - if s.active_sub is not None and s.active_sub.turns: - s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) - if s.main.turns: - s.segments.append(make_turn_segment(s.main.turns, kind="final")) + async def _respond(self, request, body, reply, in_tok, out_tok, stream) -> web.StreamResponse: + blocks, stop_reason = reply.wire + if stream: + return await _render_stream(request, blocks, stop_reason, in_tok, out_tok) + return web.json_response(_render_response(body, blocks, stop_reason, in_tok, out_tok)) - return merge_turn_segments(s.segments) +# --- Translation (Anthropic wire -> chat-template messages) --- -# ============================================================================= -# 2. Per-turn stages -# ============================================================================= - - -def _select_chain(s: Session, body: dict) -> tuple[Chain, bool, str]: - """Decide which chain this turn operates on. - - 1. fingerprint body.messages and body.system into hashes - 2. if main now contains the tool_result for a pending sub dispatch, - snapshot the sub chain into s.segments and clear s.active_sub - 3. pick main vs s.active_sub based on whether request continues main's prefix - 4. classify as 'new' | 'append' | 'wipe' against the chosen target; - a wipe also snapshots the target's current state into s.segments - - Returns (target_chain, is_sub, kind). - """ - all_msgs = body.get("messages") or [] - msg_hashes = [_hash(m) for m in all_msgs] - req_system_hash = _hash(body.get("system")) if "system" in body else s.main.system_hash - - # Close active sub-agent if its dispatch tool_result has landed on main. - if s.pending_dispatch_id and s.active_sub is not None: - tu_id = s.pending_dispatch_id - for m in all_msgs: - if not isinstance(m, dict) or m.get("role") != "user": - continue - content = m.get("content") - if not isinstance(content, list): - continue - done = any( - isinstance(b, dict) and b.get("type") == "tool_result" and b.get("tool_use_id") == tu_id - for b in content - ) - if done: - if s.active_sub.turns: - s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) - s.active_sub = None - s.pending_dispatch_id = "" - break - # Route: main iff request continues main's prefix. Sub system_hash can be - # "" (armed before sub dialled in), so never route by sub equality alone. - if s.active_sub is None: - target, is_sub = s.main, False - else: - main_continues = ( - req_system_hash == s.main.system_hash - and len(msg_hashes) >= s.main.seen_msgs - and msg_hashes[: s.main.seen_msgs] == s.main.msg_hashes[: s.main.seen_msgs] - ) - target, is_sub = (s.main, False) if main_continues else (s.active_sub, True) - - # Classify; snapshot a "wipe" segment first if we're discarding work. - if target.seen_msgs == 0: - kind = "new" - else: - is_append = ( - req_system_hash == target.system_hash - and len(msg_hashes) >= target.seen_msgs - and msg_hashes[: target.seen_msgs] == target.msg_hashes[: target.seen_msgs] - ) - if is_append: - kind = "append" - else: - if target.turns: - s.segments.append(make_turn_segment(target.turns, kind="wipe")) - kind = "wipe" - - return target, is_sub, kind - - -def _flatten(c: Any) -> str: - """Recursive Anthropic content flattener: text/tool_result(content) joined - by newline, images replaced with a placeholder.""" - if c is None: - return "" - if isinstance(c, str): - return c - if not isinstance(c, list): - return str(c) - parts: list[str] = [] - for b in c: - if isinstance(b, dict): - t = b.get("type") - if t == "text": - parts.append(b.get("text", "")) - elif t == "tool_result": - parts.append(_flatten(b.get("content"))) - elif t == "image": - parts.append("[image omitted]") - elif isinstance(b, str): - parts.append(b) - return "\n".join(p for p in parts if p) - - -def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: +def _translate_messages(msgs: list[dict], system: Any) -> list[dict]: """Anthropic messages + system -> chat-template messages. Pure function.""" translated: list[dict] = [] if system: - translated.append({"role": "system", "content": _flatten(system)}) + translated.append({"role": "system", "content": flatten_content(system)}) for m in msgs: if not isinstance(m, dict): continue role, content = m.get("role"), m.get("content") if role == "user": - blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + blocks = content if isinstance(content, list) else [{"type": "text", "text": flatten_content(content)}] for b in blocks: if isinstance(b, dict) and b.get("type") == "tool_result": - translated.append({"role": "tool", "content": _flatten(b.get("content"))}) + translated.append({"role": "tool", "content": flatten_content(b.get("content"))}) elif isinstance(b, dict) and b.get("type") == "text": translated.append({"role": "user", "content": b.get("text", "")}) else: - translated.append({"role": "user", "content": _flatten(b)}) + translated.append({"role": "user", "content": flatten_content(b)}) elif role == "assistant": texts, thinkings, tcs = [], [], [] - blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + blocks = content if isinstance(content, list) else [{"type": "text", "text": flatten_content(content)}] for b in blocks: if not isinstance(b, dict): continue @@ -208,7 +108,8 @@ def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: elif b.get("type") == "thinking": thinkings.append(b.get("thinking", "")) elif b.get("type") == "tool_use": - tcs.append({"function": {"name": b.get("name", "tool"), "arguments": b.get("input") or {}}}) + # drop the wire-only id; tool_call_dict keeps arguments a dict + tcs.append(tool_call_dict(b.get("name", "tool"), b.get("input"))) mo: dict[str, Any] = {"role": "assistant", "content": "".join(texts)} if thinkings: mo["reasoning_content"] = "".join(thinkings) @@ -216,11 +117,11 @@ def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: mo["tool_calls"] = tcs translated.append(mo) elif role == "system": - translated.append({"role": "system", "content": _flatten(content)}) + translated.append({"role": "system", "content": flatten_content(content)}) return translated -def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] | None: +def _tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] | None: """Convert Anthropic tools to tokenizer chat-template tool schema.""" if not anth_tools: return None @@ -241,157 +142,61 @@ def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] return ts or None -def _replace_chat_messages(target: Chain, body: dict) -> None: - """new/wipe: full reset of chat state and turn log.""" - all_msgs = body.get("messages") or [] - target.chat_messages = _translate_anthropic(all_msgs, body.get("system")) - if "system" in body: - target.system_hash = _hash(body.get("system")) - target.turns.clear() - target.seen_msgs = len(all_msgs) - target.msg_hashes = [_hash(m) for m in all_msgs] - if target.tools_schema is None: - target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) - - -def _extend_chat_messages(target: Chain, body: dict) -> None: - """append: translate only the new tail.""" - all_msgs = body.get("messages") or [] - translated = _translate_anthropic(all_msgs[target.seen_msgs :], None) - target.chat_messages.extend(translated) - - target.seen_msgs = len(all_msgs) - target.msg_hashes = [_hash(m) for m in all_msgs] - if target.tools_schema is None: - target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) - - -def _build_prompt(target: Chain, body: dict, kind: str, tok) -> list[int]: - """Replace/extend chat_messages and render input ids for vLLM.""" - (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, body) - return render_token_ids(target, tok) - +# --- Reply building: parsed output -> Anthropic blocks + manager_message --- -async def _generate( - prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None -) -> TurnRecord: - """Call vLLM and return a TurnRecord. - - 1. build sampling_params (session defaults overlaid with body overrides) - 2. POST vLLM ``/inference/v1/generate``; on cancel/error tear down - the request (vLLM has no per-request HTTP abort endpoint) - 3. keep the exact prompt/output token ids; trajectory merge later compares - later prompt tokens with earlier outputs to build the loss mask - """ - return await call_vllm_generate( - prompt_ids, - s, - body, - app, - max_token_keys=("max_tokens",), - stop_keys=("stop_sequences",), - log_prefix="anthropic_adapter", - logger=logger, - session_id=session_id, - ) +def _build_reply_parts( + parsed: ParsedModelOutput, + finish: str, +) -> tuple[list[dict], str, dict[str, Any]]: + """Return (anthropic blocks, wire stop_reason, manager_message). -def _build_reply(target: Chain, output_ids: list[int], finish: str, app) -> tuple[list[dict], str, str]: - """Turn the model's raw output ids into the reply we send back to claude-code. - - 1. parse decoded text -> (thinking, visible, tool_uses) via parsers - 2. pack into Anthropic content blocks; tag dispatch_id when a tool_use - names Task/Agent (sub-agent trigger) - 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' - - Returns (blocks, stop_reason, dispatch_id). + The tool_calls inside manager_message use canonical args (tool_call_dict) so + this assistant turn compares equal (dict equality) to the same turn replayed + as history on the next request. """ - tok = app[TOKENIZER_KEY] - - raw_output = tok.decode(output_ids, skip_special_tokens=False) if output_ids else "" - parsed = parse_model_output( - raw_output, - tokenizer=tok, - tools_schema=target.tools_schema, - tool_parser_name=app[TOOL_PARSER_KEY], - reasoning_parser_name=app[REASONING_PARSER_KEY], - ) - blocks, dispatch_id = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) - return blocks, _stop_reason(parsed.tool_uses, finish), dispatch_id - - -def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> tuple[list[dict], str]: - """Pack parsed model output into Anthropic content blocks.""" blocks: list[dict] = [] - if thinking: - blocks.append({"type": "thinking", "thinking": thinking}) - if visible: - blocks.append({"type": "text", "text": visible}) - dispatch_id = "" - for tu in tool_uses: + if parsed.reasoning: + blocks.append({"type": "thinking", "thinking": parsed.reasoning}) + if parsed.text: + blocks.append({"type": "text", "text": parsed.text}) + + manager_tcs: list[dict] = [] + for tu in parsed.tool_uses: tu_id = f"toolu_{secrets.token_hex(8)}" blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) - if tu["name"] in _SUBAGENT_TOOLS: - dispatch_id = tu_id + # tu_id is wire-only; tool_call_dict drops it so the leaf matches its echo + manager_tcs.append(tool_call_dict(tu["name"], tu.get("input"))) + if not blocks: blocks.append({"type": "text", "text": ""}) - return blocks, dispatch_id - -def _stop_reason(tool_uses: list[dict], finish: str) -> str: - if tool_uses: - return "tool_use" - if finish == "length": - return "max_tokens" - return "end_turn" + if parsed.tool_uses: + stop_reason = "tool_use" + elif finish == "length": + stop_reason = "max_tokens" + else: + stop_reason = "end_turn" + manager_message: dict[str, Any] = {"role": "assistant", "content": parsed.text or ""} + if parsed.reasoning: + manager_message["reasoning_content"] = parsed.reasoning + if manager_tcs: + manager_message["tool_calls"] = manager_tcs -def _start_sub_chain(s: Session, dispatch_id: str) -> None: - """Start a fresh sub chain on this session and remember the tool_use_id - we'll watch for on main to know when this sub is done. The matching - 'sub done' step lives inside _select_chain.""" - s.pending_dispatch_id = dispatch_id - if s.active_sub is None: - s.active_sub = Chain() + return blocks, stop_reason, manager_message -# ============================================================================= -# 3. Request handling -- one full turn + SSE wrap -# ============================================================================= +# --- Request framing: session id + wire response/stream rendering --- def _request_session_id(request: web.Request) -> str: - return request_session_id(request, include_x_api_key=True) - - -async def _handle_request(request: web.Request) -> web.StreamResponse: - body = await request.json() - sid = _request_session_id(request) - adapter = request.app[ADAPTER_KEY] - if sid in adapter.closed: # session drained; refuse stragglers - return web.Response(status=503, text="session closed") - app = request.app - s = adapter.store.setdefault(sid, Session()) - task = asyncio.current_task() - adapter.inflight.setdefault(sid, set()).add(task) - try: - async with s.lock: # same sid -> serialized - target, is_sub, kind = _select_chain(s, body) - ideal_ids = _build_prompt(target, body, kind, app[TOKENIZER_KEY]) - turn = await _generate(ideal_ids, s, body, app, session_id=sid) - blocks, stop, did = _build_reply(target, turn.output_ids, turn.finish_reason, app) - target.turns.append(turn) - if did and not is_sub: # sub doesn't nest - _start_sub_chain(s, did) - in_tok, out_tok = len(ideal_ids), len(turn.output_ids) - if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): - return await _stream_response(request, blocks, stop, in_tok, out_tok) - return web.json_response(_message_response(body, blocks, stop, in_tok, out_tok)) - finally: - adapter.inflight.get(sid, set()).discard(task) - - -def _message_response(body: dict, blocks: list[dict], stop_reason: str, in_tok: int, out_tok: int) -> dict: + # Anthropic auth lands in Authorization: Bearer or X-Api-Key; the Messages + # body carries no sid hint. Bearer wins when both are present. + return sid_from_bearer(request) or (request.headers.get("X-Api-Key") or "").strip() or "default" + + +def _render_response(body: dict, blocks: list[dict], stop_reason: str, in_tok: int, out_tok: int) -> dict: return { "id": f"msg_{secrets.token_hex(12)}", "type": "message", @@ -404,10 +209,10 @@ def _message_response(body: dict, blocks: list[dict], stop_reason: str, in_tok: } -async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web.StreamResponse: - """Stream blocks back to claude-code as an Anthropic Messages SSE - response: message_start, (content_block_start, content_block_delta, - content_block_stop)*N, message_delta, message_stop.""" +async def _render_stream(request, blocks, stop_reason, in_tok, out_tok) -> web.StreamResponse: + """Stream blocks back as an Anthropic Messages SSE response: message_start, + (content_block_start, content_block_delta, content_block_stop)*N, + message_delta, message_stop.""" out = web.StreamResponse( status=200, headers={ @@ -418,7 +223,6 @@ async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web ) await out.prepare(request) - # message_start ms_data = { "type": "message_start", "message": { @@ -471,12 +275,77 @@ async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web return out -# Trivial endpoints claude-code probes during a session: count_tokens runs -# every turn (return 0 -- client uses it as a hint, not a hard budget), -# healthz/v1/models are startup readiness checks. +# count_tokens runs every turn but the client uses it only as a hint, not a +# hard budget, so returning 0 is fine. async def _count_tokens(request: web.Request) -> web.Response: + await request.read() return web.json_response({"input_tokens": 0}) -async def _ok(request: web.Request) -> web.Response: - return await ok_response(request) +# --- Anthropic-specific quirks: mid-list system folding --- + + +_MID_SYSTEM_WRAP_PREFIX = "\n" +_MID_SYSTEM_WRAP_SUFFIX = "\n\n" + + +def _fold_mid_list_system_into_user(body_obj: dict) -> bool: + """Fold non-leading role:system messages into a neighbouring user message as + a text block. Mutates body_obj in place; returns True iff + any fold happened. + + Some clients insert a system message in the middle of the message list, but + many chat templates reject any system message past index 0. Attaching the + wrapped reminder to the preceding user message (or the next one, if there is + no prior user message) keeps the history acceptable to the template. + """ + msgs = body_obj.get("messages") + if not isinstance(msgs, list) or not msgs: + return False + + system_idx = [i for i, m in enumerate(msgs) if isinstance(m, dict) and m.get("role") == "system" and i > 0] + if not system_idx: + return False + + def _promote_to_list(msg: dict) -> list: + c = msg.get("content") + if isinstance(c, list): + return c + msg["content"] = [{"type": "text", "text": c if isinstance(c, str) else ""}] + return msg["content"] + + def _wrap(text: str) -> dict: + return { + "type": "text", + "text": _MID_SYSTEM_WRAP_PREFIX + text + _MID_SYSTEM_WRAP_SUFFIX, + } + + changed = False + TOMBSTONE: dict = {"__folded__": True} + for i in system_idx: + sys_msg = msgs[i] + wrapped = _wrap(flatten_content(sys_msg.get("content"))) + target = None + for j in range(i - 1, -1, -1): + cand = msgs[j] + if isinstance(cand, dict) and cand.get("role") == "user": + target = cand + _promote_to_list(target).append(wrapped) + break + if target is None: + for j in range(i + 1, len(msgs)): + cand = msgs[j] + if isinstance(cand, dict) and cand.get("role") == "user": + target = cand + _promote_to_list(target).insert(0, wrapped) + break + if target is None: + msgs[i] = {"role": "user", "content": [wrapped]} + changed = True + continue + msgs[i] = TOMBSTONE + changed = True + + if changed: + body_obj["messages"] = [m for m in msgs if m is not TOMBSTONE] + return changed diff --git a/vime/agent/adapters/common.py b/vime/agent/adapters/common.py index c598be9f1..15899d067 100644 --- a/vime/agent/adapters/common.py +++ b/vime/agent/adapters/common.py @@ -1,150 +1,416 @@ -"""Shared adapter primitives for token-capturing agent rollouts.""" +"""Shared adapter primitives for token-capturing agent rollouts. + +A protocol adapter (Anthropic / OpenAI) subclasses BaseAdapter and fills in the +wire-specific hooks (_register_routes, _session_id, _translate, _build_reply, +_respond, and optionally _preprocess_body) plus a few class attributes (logger, +log_prefix, max_token_keys, stop_keys). The session lifecycle, per-sid turn cap, +inflight-task bookkeeping and the one-turn _run_turn pipeline are inherited. + +flatten_content, tool_call_dict and manager_finish_reason cover the parts both +protocols handle identically. +""" from __future__ import annotations import asyncio import dataclasses -import hashlib -import json import logging +import time from collections.abc import Callable from typing import Any import aiohttp from aiohttp import web -from vime.agent.trajectory import TokenSegment, TurnRecord +from vime.agent.parsing import parse_model_output +from vime.agent.trajectory import TrajectoryManager, TurnRecord -ADAPTER_KEY = web.AppKey("adapter", object) -TOKENIZER_KEY = web.AppKey("tokenizer", object) -VLLM_URL_KEY = web.AppKey("vllm_url", object) -TOOL_PARSER_KEY = web.AppKey("tool_parser", object) -REASONING_PARSER_KEY = web.AppKey("reasoning_parser", object) +__all__ = ["TurnRecord"] @dataclasses.dataclass -class AdapterChain: - """Protocol-neutral chat chain state used by HTTP adapters.""" +class Session: + """Per-sid adapter state: sampling defaults and context budget. + + Trajectory state lives in the shared TrajectoryManager (BaseAdapter.manager), + not here. + """ + + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + + +@dataclasses.dataclass +class Reply: + """Output of an adapter's _build_reply, consumed by _run_turn. + + manager_message and finish_reason feed record_turn and the debug callback; + wire is opaque to the pipeline and only the adapter's own _respond reads it. + """ + + manager_message: dict + finish_reason: str + wire: Any + + +def _render_token_ids( + messages: list[dict], + tokenizer, + *, + tools: list[dict] | None, + add_generation_prompt: bool = True, +) -> list[int]: + """Render a chat-message list to token ids with the served chat template.""" + enc = tokenizer.apply_chat_template( + messages, + tools=tools, + tokenize=True, + add_generation_prompt=add_generation_prompt, + ) + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) + - system_hash: str = "" - chat_messages: list[dict] = dataclasses.field(default_factory=list) - tools_schema: list[dict] | None = None - seen_msgs: int = 0 - msg_hashes: list[str] = dataclasses.field(default_factory=list) - turns: list[TurnRecord] = dataclasses.field(default_factory=list) +def flatten_content(c: Any) -> str: + """Flatten a wire content value into a chat-template string. + + Handles both Anthropic and OpenAI block shapes. A non-list value (str / + dict / other) is returned via str() unchanged. + """ + if c is None: + return "" + if isinstance(c, str): + return c + if not isinstance(c, list): + return str(c) + parts: list[str] = [] + for b in c: + if isinstance(b, str): + parts.append(b) + continue + if not isinstance(b, dict): + parts.append(str(b)) + continue + t = b.get("type") + if t in {"text", "input_text", "output_text"}: + parts.append(b.get("text", "")) + elif t == "tool_result": + parts.append(flatten_content(b.get("content"))) + elif t in {"image", "image_url", "input_image"}: + parts.append("[image omitted]") + elif "content" in b: + parts.append(flatten_content(b.get("content"))) + elif "text" in b: + parts.append(str(b.get("text") or "")) + return "\n".join(p for p in parts if p) + + +def tool_call_dict(name: str, arguments: dict | None) -> dict: + """Canonical OpenAI-shape tool call stored on manager_message. + + arguments stays a dict (not a JSON string): the chat template needs a + mapping, and the trajectory manager matches history by dict equality, so a + sampled leaf and its replayed echo compare equal regardless of key order. + The wire-only tool-call id is dropped for the same reason. + """ + return {"type": "function", "function": {"name": name, "arguments": arguments or {}}} + + +def manager_finish_reason(tool_uses: list[dict], raw_finish: str) -> str: + """Finish reason stored on the manager turn: tool_calls if the turn called a + tool, else the raw vllm finish.""" + return "tool_calls" if tool_uses else (raw_finish or "stop") class BaseAdapter: - """Base HTTP adapter with per-instance session lifecycle state.""" + """Base HTTP adapter: session lifecycle plus the shared one-turn pipeline. - session_cls: type + See the module docstring for the class attributes and hooks a subclass must + supply; everything else is inherited. + """ - def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + logger: logging.Logger = logging.getLogger(__name__) + log_prefix: str = "adapter" + # body keys that cap max_new_tokens and carry stop sequences, in priority order + max_token_keys: tuple[str, ...] = () + stop_keys: tuple[str, ...] = () + manager: Any + + def __init__( + self, + *, + tokenizer, + vllm_url, + tool_parser=None, + reasoning_parser=None, + max_turns_per_sid: int | None = None, + fork_threshold_tokens: int | None = None, + debug_callback: Callable[..., None] | None = None, + ) -> None: + self.tokenizer = tokenizer + self.vllm_url = vllm_url.rstrip("/") if isinstance(vllm_url, str) else vllm_url + self.tool_parser = tool_parser + self.reasoning_parser = reasoning_parser self.store: dict[str, Any] = {} self.inflight: dict[str, set[asyncio.Task]] = {} self.closed: set[str] = set() self.app = web.Application(client_max_size=64 * 1024 * 1024) - self.app[ADAPTER_KEY] = self - self.app[TOKENIZER_KEY] = tokenizer - self.app[VLLM_URL_KEY] = vllm_url.rstrip("/") if isinstance(vllm_url, str) else vllm_url - self.app[TOOL_PARSER_KEY] = tool_parser - self.app[REASONING_PARSER_KEY] = reasoning_parser - def open_session( - self, - sid: str, - *, - sampling_defaults: dict | None = None, - max_context_tokens: int = 0, - ) -> None: - register_session( - self.store, - sid, - self.session_cls, - sampling_defaults=sampling_defaults, - max_context_tokens=max_context_tokens, - ) + # one manager shared across all sids; per-sid trees live inside it. + # fork_threshold_tokens left None means the manager uses its own default. + mgr_kwargs: dict[str, int] = {} + if fork_threshold_tokens is not None: + mgr_kwargs["fork_threshold_tokens"] = fork_threshold_tokens + self.manager = TrajectoryManager(**mgr_kwargs) - async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: - await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) + self.debug_callback: Callable[..., None] | None = debug_callback + # per-sid turn cap: return 429 to kill the run once exceeded + self.max_turns_per_sid: int | None = max_turns_per_sid + self._sid_turn_count: dict[str, int] = {} - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - raise NotImplementedError + self.app.router.add_get("/healthz", _health) + self.app.router.add_get("/v1/models", _health) + self._register_routes(self.app) + # -- wire hooks (subclass overrides) ------------------------------------- -def strip_cache_control(obj: Any) -> Any: - if isinstance(obj, dict): - return {k: strip_cache_control(v) for k, v in obj.items() if k != "cache_control"} - if isinstance(obj, list): - return [strip_cache_control(x) for x in obj] - return obj + def _register_routes(self, app: web.Application) -> None: + """Register the protocol's POST route(s) and bind self._run_turn.""" + raise NotImplementedError + def _session_id(self, request: web.Request, body: dict) -> str: + raise NotImplementedError -def stable_hash(obj: Any) -> str: - payload = json.dumps(strip_cache_control(obj), sort_keys=True, ensure_ascii=False, default=str).encode("utf-8") - return hashlib.sha1(payload).hexdigest()[:12] + def _preprocess_body(self, body: dict) -> None: + """Mutate the parsed body in place before sid resolution (default no-op).""" + def _translate(self, body: dict) -> tuple[list[dict], list[dict] | None]: + """Return (chat_messages, tools_schema) from the wire body.""" + raise NotImplementedError -def json_arguments(value: Any) -> str: - if value is None: - return "{}" - if isinstance(value, str): - return value - return json.dumps(value, ensure_ascii=False) + def _build_reply(self, parsed, raw_finish: str, translated: list[dict], tools_schema: list[dict] | None) -> Reply: + """Pack parsed model output into a Reply.""" + raise NotImplementedError + async def _respond( + self, + request: web.Request, + body: dict, + reply: Reply, + in_tok: int, + out_tok: int, + stream: bool, + ) -> web.StreamResponse: + raise NotImplementedError -def render_token_ids(chain: AdapterChain, tokenizer) -> list[int]: - enc = tokenizer.apply_chat_template( - chain.chat_messages, - tools=chain.tools_schema, - tokenize=True, - add_generation_prompt=True, - ) - ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc - return list(ids) + # -- session lifecycle --------------------------------------------------- + def open_session( + self, + sid: str, + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, + ) -> None: + """Register a fresh per-sid Session; sids must be unique.""" + if sid in self.store: + raise ValueError(f"session_id {sid!r} already exists; sids must be unique per agent run") + self.store[sid] = Session( + sampling_defaults=dict(sampling_defaults or {}), + max_context_tokens=int(max_context_tokens or 0), + ) -def request_session_id( - request: web.Request, - *, - body: dict | None = None, - include_x_api_key: bool = False, -) -> str: - auth = request.headers.get("Authorization", "") - if auth.lower().startswith("bearer "): - sid = auth[7:].strip() - if sid: - return sid + async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: + """Mark a sid closed and drain its in-flight turn tasks.""" + self.closed.add(sid) + tasks = [t for t in self.inflight.pop(sid, ()) if not t.done()] + if not tasks: + return + + async def _drain() -> None: + _, pending = await asyncio.wait(tasks, timeout=wait_timeout) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + loop = tasks[0].get_loop() + try: + await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(_drain(), loop)) + except Exception: + self.logger.exception("[%s] sid=%s shutdown drain failed", self.log_prefix, sid) + + async def finish_session( + self, + sid: str, + *, + base_sample, + reward: float = 0.0, + extra_metadata: dict | None = None, + wait_timeout: float = 5.0, + ) -> list: + """Drain a session's trajectory into fully-formed Sample objects. + + Waits out in-flight requests for the sid, linearises the per-sid tree, + then decodes each sample's trained tail into .response (the manager is + tokenizer-free, so the adapter that owns the tokenizer fills this in). + Idempotent: a second call for an already-popped sid returns []. + """ + await self.shutdown_session(sid, wait_timeout=wait_timeout) + session = self.store.pop(sid, None) + max_sample_tokens = int(getattr(session, "max_context_tokens", 0) or 0) if session is not None else 0 + samples = self.manager.get_trajectory( + sid, + base_sample=base_sample, + reward=reward, + extra_metadata=extra_metadata, + max_sample_tokens=max_sample_tokens, + ) + for s in samples: + rlen = int(s.response_length or 0) + s.response = ( + self.tokenizer.decode(s.tokens[-rlen:], skip_special_tokens=False) if rlen and s.tokens else "" + ) + return samples + + async def drop_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: + await self.shutdown_session(sid, wait_timeout=wait_timeout) + self.store.pop(sid, None) + self.manager.drop_session(sid) + + # -- shared request pipeline --------------------------------------------- + + def _check_turn_cap(self, sid: str) -> web.Response | None: + """Enforce max_turns_per_sid, returning a 429 response once exceeded. + + Increments the per-sid counter as a side effect when under the cap. + """ + cap = self.max_turns_per_sid + if cap is None: + return None + prior = self._sid_turn_count.get(sid, 0) + if prior >= cap: + self.logger.warning("[%s] sid=%s exceeded max_turns_per_sid=%d; killing run", self.log_prefix, sid, cap) + return web.json_response( + { + "error": { + "type": "rate_limit_error", + "message": (f"adapter: sid {sid!r} exceeded max_turns_per_sid={cap}; killing run"), + } + }, + status=429, + ) + self._sid_turn_count[sid] = prior + 1 + return None + + def _run_debug_callback(self, sid, translated, tools_schema, manager_message, turn) -> None: + """Run the optional debug-only data dump callback; unset in production.""" + callback = self.debug_callback + if callback is None: + return + try: + callback(sid, translated, tools_schema, manager_message, turn) + except Exception: + self.logger.exception("debug_callback failed (sid=%s)", sid) + + async def _run_turn(self, request: web.Request) -> web.StreamResponse: + """One full agent turn: translate -> vllm -> parse -> append -> respond. + + The wire-specific steps are delegated to the subclass hooks; the rest + (sid resolution, closed/cap guards, inflight tracking, record_turn) is + shared across protocols. + """ + body = await request.json() + self._preprocess_body(body) + sid = self._session_id(request, body) + if sid in self.closed: # session drained; refuse stragglers + self.logger.debug("[%s] sid=%s request after session closed", self.log_prefix, sid) + return web.Response(status=503, text="session closed") + capped = self._check_turn_cap(sid) + if capped is not None: + return capped + + tok = self.tokenizer + s = self.store.setdefault(sid, Session()) + task = asyncio.current_task() + self.inflight.setdefault(sid, set()).add(task) + t0 = time.monotonic() + try: + translated, tools_schema = self._translate(body) + prompt_ids = _render_token_ids(translated, tok, tools=tools_schema, add_generation_prompt=True) + + turn = await call_vllm_generate(prompt_ids, s, body, adapter=self, session_id=sid) + + raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" + parsed = parse_model_output( + raw_output, + tokenizer=tok, + tools_schema=tools_schema, + tool_parser_name=self.tool_parser, + reasoning_parser_name=self.reasoning_parser, + ) + reply = self._build_reply(parsed, turn.finish_reason, translated, tools_schema) + turn = dataclasses.replace(turn, ill_formed=parsed.ill_formed) + + in_tok, out_tok = len(prompt_ids), len(turn.output_ids) + stream = body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", "") + + # Flush the response before recording the trajectory: a client that + # disconnected during generation makes _respond raise here, and we + # must not record a turn the client never received. + try: + response = await self._respond(request, body, reply, in_tok, out_tok, stream) + except (ConnectionResetError, asyncio.CancelledError) as e: + self.logger.warning( + "[%s] sid=%s client disconnected before response flush: %s after %.1fs", + self.log_prefix, + sid, + type(e).__name__, + time.monotonic() - t0, + ) + if isinstance(e, asyncio.CancelledError): + raise + return web.Response(status=499, text="client disconnected") + + self._run_debug_callback( + sid, + translated, + tools_schema, + reply.manager_message, + turn, + ) - if body is not None: - metadata = body.get("metadata") - if isinstance(metadata, dict) and metadata.get("session_id"): - return str(metadata["session_id"]) - if body.get("user"): - return str(body["user"]) + self.manager.record_turn( + sid, + turn=turn, + prompt_messages=translated, + response_message=reply.manager_message, + metadata={"sid": sid}, + ) + return response + finally: + self.inflight.get(sid, set()).discard(task) - if include_x_api_key: - api_key = request.headers.get("X-Api-Key") - if api_key: - return api_key.strip() - return "default" +def sid_from_bearer(request: web.Request) -> str | None: + """sid from the Authorization: Bearer header, or None if absent.""" + auth = request.headers.get("Authorization", "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() or None + return None -def register_session( - store: dict[str, Any], - sid: str, - session_factory: Callable[[], Any], - *, - sampling_defaults: dict | None = None, - max_context_tokens: int = 0, -) -> None: - if sid in store: - raise ValueError(f"session_id {sid!r} already exists; sids must be unique per agent run") - session = store[sid] = session_factory() - session.sampling_defaults = dict(sampling_defaults or {}) - session.max_context_tokens = int(max_context_tokens or 0) +def sid_from_body(body: dict | None) -> str | None: + """sid from the OpenAI-shape body (metadata.session_id / user), or None.""" + if not body: + return None + metadata = body.get("metadata") + if isinstance(metadata, dict) and metadata.get("session_id"): + return str(metadata["session_id"]) + if body.get("user"): + return str(body["user"]) + return None def _sampling_params(session: Any, body: dict, *, max_token_keys: tuple[str, ...], stop_keys: tuple[str, ...]) -> dict: @@ -223,29 +489,31 @@ async def call_vllm_generate( prompt_ids: list[int], session: Any, body: dict, - app, *, - max_token_keys: tuple[str, ...], - stop_keys: tuple[str, ...], - log_prefix: str, - logger: logging.Logger, + adapter: BaseAdapter, session_id: str | None = None, ) -> TurnRecord: - sp = _sampling_params(session, body, max_token_keys=max_token_keys, stop_keys=stop_keys) + """POST one turn to vllm ``/inference/v1/generate`` and pack the reply into a TurnRecord. + + Module-level (not a method) so tests can monkeypatch it. + """ + logger = adapter.logger + sp = _sampling_params(session, body, max_token_keys=adapter.max_token_keys, stop_keys=adapter.stop_keys) if session.max_context_tokens > 0: remaining_context = session.max_context_tokens - len(prompt_ids) if remaining_context <= 0: logger.warning( - "[%s] prompt exceeds max_context_tokens (%d >= %d)", - log_prefix, + "[%s] sid=%s prompt exceeds max_context_tokens (%d >= %d)", + adapter.log_prefix, + session_id, len(prompt_ids), session.max_context_tokens, ) return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[], finish_reason="length") sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", remaining_context)), remaining_context) - vllm_url = app[VLLM_URL_KEY] + vllm_url = adapter.vllm_url payload: dict[str, Any] = { "token_ids": list(prompt_ids), "sampling_params": _vllm_sampling_body(sp), @@ -263,16 +531,24 @@ async def call_vllm_generate( ) as r: if r.status >= 400: text = await r.text() + logger.warning( + "[%s] sid=%s vllm upstream %d: %.200s", + adapter.log_prefix, + session_id, + r.status, + text, + ) raise RuntimeError(f"vllm upstream {r.status}: {text[:400]}") data = await r.json(content_type=None) choice = (data.get("choices") or [{}])[0] output_ids, output_log_probs = _tokens_and_logprobs_from_choice(choice) fr = choice.get("finish_reason") finish = fr if isinstance(fr, str) and fr else "stop" - except (asyncio.CancelledError, aiohttp.ClientError, asyncio.TimeoutError): + except (asyncio.CancelledError, aiohttp.ClientError, asyncio.TimeoutError) as e: # vLLM ``/inference/v1/generate`` has no per-request HTTP abort endpoint. # Cancelling the in-flight task tears down the aiohttp request, which drops # the streaming connection so vLLM stops generating. + logger.debug("[%s] sid=%s turn aborted: %s", adapter.log_prefix, session_id, type(e).__name__) if task is not None: task.cancel() raise @@ -285,23 +561,6 @@ async def call_vllm_generate( ) -async def shutdown_session_tasks( - sid: str, - closed: set[str], - inflight: dict[str, set[asyncio.Task]], - *, - wait_timeout: float = 5.0, -) -> None: - closed.add(sid) - tasks = [t for t in inflight.pop(sid, ()) if not t.done()] - if not tasks: - return - _, pending = await asyncio.wait(tasks, timeout=wait_timeout) - for task in pending: - task.cancel() - if pending: - await asyncio.gather(*pending, return_exceptions=True) - - -async def ok_response(request: web.Request) -> web.Response: +async def _health(request: web.Request) -> web.Response: + """Handler for /healthz and /v1/models readiness probes.""" return web.json_response({"ok": True}) diff --git a/vime/agent/adapters/openai.py b/vime/agent/adapters/openai.py index 89ca06707..62d465fa2 100644 --- a/vime/agent/adapters/openai.py +++ b/vime/agent/adapters/openai.py @@ -1,17 +1,19 @@ -"""OpenAI-compatible adapters for agent rollouts. - -The adapter exposes ``/v1/chat/completions`` and ``/v1/responses``. Both -endpoints render incoming messages with the served model's chat template, call -vLLM's ``/inference/v1/generate`` with ``token_ids``, and record -the exact sampled token ids/logprobs as ``TurnRecord`` objects. New code should -use ``OpenAIAdapter`` and call ``finish_session()`` at trajectory end to drain -trainable ``TokenSegment`` objects. +"""OpenAI Chat-Completions adapter for agent rollouts. + +Mirrors vime.agent.adapters.anthropic but speaks the OpenAI +/v1/chat/completions protocol, so an OpenAI-compatible client (e.g. the Codex +CLI) can drive the vime vllm server. Each request is rendered with the served +model's chat template, sent to vllm ``/inference/v1/generate`` as ``token_ids``, +parsed, and folded into a shared TrajectoryManager keyed by session id. +finish_session(sid) drains a session's trajectory into a list of Sample. + +Only /v1/chat/completions is implemented; the Responses API (/v1/responses) is +out of scope. The section layout (adapter class -> translation -> reply building +-> request framing) mirrors vime.agent.adapters.anthropic. """ from __future__ import annotations -import asyncio -import dataclasses import json import logging import secrets @@ -20,328 +22,272 @@ from aiohttp import web -from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY -from vime.agent.adapters.common import AdapterChain as Chain -from vime.agent.adapters.common import BaseAdapter, call_vllm_generate -from vime.agent.adapters.common import json_arguments as _json_arguments -from vime.agent.adapters.common import ok_response, render_token_ids, request_session_id -from vime.agent.adapters.common import stable_hash as _hash -from vime.agent.parsing import ParsedModelOutput, parse_model_output -from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments +from vime.agent.adapters.common import ( + BaseAdapter, + Reply, + flatten_content, + manager_finish_reason, + sid_from_bearer, + sid_from_body, +) +from vime.agent.parsing import ParsedModelOutput logger = logging.getLogger(__name__) -@dataclasses.dataclass -class Session: - main: Chain = dataclasses.field(default_factory=Chain) - sampling_defaults: dict = dataclasses.field(default_factory=dict) - max_context_tokens: int = 0 - lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - segments: list[TurnSegment] = dataclasses.field(default_factory=list) +class OpenAIAdapter(BaseAdapter): + """OpenAI Chat-Completions-compatible HTTP adapter: wire translation and + reply framing only; the turn machinery is inherited from BaseAdapter.""" + + logger = logger + log_prefix = "openai_adapter" + max_token_keys = ("max_completion_tokens", "max_tokens", "max_output_tokens") + stop_keys = ("stop",) + + def _register_routes(self, app: web.Application) -> None: + app.router.add_post("/v1/chat/completions", self._run_turn) + + def _session_id(self, request: web.Request, body: dict) -> str: + return _request_session_id(request, body) + + def _translate(self, body: dict) -> tuple[list[dict], list[dict] | None]: + messages = body.get("messages") or [] + if not isinstance(messages, list): + raise web.HTTPBadRequest(text="messages must be a list") + translated = _translate_messages(messages) + tools_schema = _tools_to_chat_tools(body.get("tools")) + return translated, tools_schema + + def _build_reply(self, parsed, raw_finish, translated, tools_schema) -> Reply: + wire_message, manager_message, wire_finish = _build_reply_parts(parsed, raw_finish) + return Reply( + manager_message=manager_message, + finish_reason=manager_finish_reason(parsed.tool_uses, raw_finish), + wire=(wire_message, wire_finish), + ) + async def _respond(self, request, body, reply, in_tok, out_tok, stream) -> web.StreamResponse: + wire_message, wire_finish = reply.wire + if stream: + return await _render_stream(request, body, wire_message, wire_finish, in_tok, out_tok) + return web.json_response(_render_response(body, wire_message, wire_finish, in_tok, out_tok)) -class OpenAIAdapter(BaseAdapter): - """OpenAI-compatible HTTP adapter with session lifecycle helpers.""" - session_cls = Session +# --- Translation (OpenAI wire -> chat-template messages) --- - def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: - super().__init__( - tokenizer=tokenizer, - vllm_url=vllm_url, - tool_parser=tool_parser, - reasoning_parser=reasoning_parser, - ) - self.app.router.add_post("/v1/chat/completions", _handle_chat_completions) - self.app.router.add_post("/v1/responses", _handle_responses) - self.app.router.add_get("/healthz", _ok) - self.app.router.add_get("/v1/models", _ok) - - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - await self.shutdown_session(sid, wait_timeout=wait_timeout) - s = self.store.pop(sid, None) - if s is None: - return [] - if s.main.turns: - s.segments.append(make_turn_segment(s.main.turns, kind="final")) - return merge_turn_segments(s.segments) - - -def _flatten_content(content: Any) -> str: - """Flatten OpenAI text/content parts into a chat-template string.""" - if content is None: - return "" - if isinstance(content, str): - return content - if not isinstance(content, list): - return str(content) - - parts: list[str] = [] - for item in content: - if isinstance(item, str): - parts.append(item) - continue - if not isinstance(item, dict): - parts.append(str(item)) - continue - typ = item.get("type") - if typ in {"text", "input_text", "output_text"}: - parts.append(item.get("text", "")) - elif typ in {"image_url", "input_image"}: - parts.append("[image omitted]") - elif "content" in item: - parts.append(_flatten_content(item.get("content"))) - elif "text" in item: - parts.append(str(item.get("text") or "")) - return "\n".join(p for p in parts if p) - - -def _normalize_tool_call(call: dict[str, Any]) -> dict[str, Any]: - function = call.get("function") or {} - name = function.get("name") or call.get("name") or "tool" - arguments = function.get("arguments", call.get("arguments", {})) - out = { - "type": "function", - "function": { - "name": name, - "arguments": _json_arguments(arguments), - }, - } - if call.get("id"): - out["id"] = call["id"] - return out +def _arguments_as_dict(arguments: Any) -> dict[str, Any]: + """Coerce wire-shape tool_calls[].function.arguments into a dict. -def _translate_chat_messages(messages: list[dict]) -> list[dict]: - """OpenAI chat messages -> tokenizer chat-template messages.""" + OpenAI sends arguments as a JSON-encoded string; the chat template and the + trajectory manager's history matching both expect a mapping. Malformed + payloads fall back to {"_raw_arguments": s}. + """ + if isinstance(arguments, dict): + return arguments + if arguments is None: + return {} + if isinstance(arguments, str): + s = arguments.strip() + if not s: + return {} + try: + parsed = json.loads(s) + except json.JSONDecodeError: + return {"_raw_arguments": arguments} + return parsed if isinstance(parsed, dict) else {"_raw_arguments": arguments} + return {"_raw_arguments": str(arguments)} + + +def _translate_messages(messages: list[dict]) -> list[dict]: + """OpenAI chat messages -> tokenizer chat-template messages. + + Mirrors anthropic._translate_messages so a replayed assistant turn compares + equal (dict equality) to the leaf the manager appended on the previous + request. Two invariants must hold: + + * tool_calls[i].function.arguments is a dict (not a JSON string): the chat + template needs a mapping, and the manager matches history by dict + equality regardless of key order. + * Wire-only correlation ids are dropped (tool_call_id on tool messages, + tool_calls[i].id on echoed assistant messages). Fresh ids are minted on + each response, so keeping the wire ids would diverge the replay match. + """ translated: list[dict] = [] for msg in messages: if not isinstance(msg, dict): continue role = msg.get("role") content = msg.get("content") - if role == "developer": + if role == "developer": # OpenAI Responses API alias role = "system" if role in {"system", "user"}: - translated.append({"role": role, "content": _flatten_content(content)}) + translated.append({"role": role, "content": flatten_content(content)}) elif role == "tool": - tool_msg = {"role": "tool", "content": _flatten_content(content)} - if msg.get("tool_call_id"): - tool_msg["tool_call_id"] = msg["tool_call_id"] - translated.append(tool_msg) + # drop tool_call_id -- wire-only correlation field; see docstring + translated.append({"role": "tool", "content": flatten_content(content)}) elif role == "assistant": - assistant: dict[str, Any] = {"role": "assistant", "content": _flatten_content(content)} - if msg.get("reasoning_content"): - assistant["reasoning_content"] = msg["reasoning_content"] + assistant: dict[str, Any] = { + "role": "assistant", + "content": flatten_content(content), + } + reasoning = msg.get("reasoning_content") + if reasoning: + assistant["reasoning_content"] = reasoning tool_calls = msg.get("tool_calls") or [] - if tool_calls: - assistant["tool_calls"] = [_normalize_tool_call(c) for c in tool_calls if isinstance(c, dict)] + normalized: list[dict[str, Any]] = [] + for call in tool_calls: + if not isinstance(call, dict): + continue + function = call.get("function") or {} + name = function.get("name") or call.get("name") or "tool" + arguments = function.get("arguments") + if arguments is None: + arguments = call.get("arguments", {}) + # NB: arguments stays a dict (not a JSON string), and the + # wire-only id is dropped. See docstring above. + normalized.append( + { + "type": "function", + "function": { + "name": name, + "arguments": _arguments_as_dict(arguments), + }, + } + ) + if normalized: + assistant["tool_calls"] = normalized translated.append(assistant) + # unknown roles are silently dropped return translated -def _normalize_tool(tool: dict[str, Any]) -> dict[str, Any] | None: - if not isinstance(tool, dict): - return None - if tool.get("type") != "function": +def _tools_to_chat_tools(tools: list[dict] | None) -> list[dict] | None: + """Convert OpenAI tools list to tokenizer chat-template tool schema.""" + if not tools: return None - if isinstance(tool.get("function"), dict): - function = tool["function"] - name = function.get("name") - if not name: - return None - return { - "type": "function", - "function": { - "name": name, - "description": function.get("description", ""), - "parameters": function.get("parameters") or {"type": "object", "properties": {}}, - }, - } - name = tool.get("name") - if not name: - return None - return { - "type": "function", - "function": { - "name": name, - "description": tool.get("description", ""), - "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, - }, - } - - -def _normalize_tools(tools: list[dict] | None) -> list[dict] | None: - normalized = [_normalize_tool(t) for t in tools or []] - return [t for t in normalized if t is not None] or None - - -def _responses_input_to_messages(input_value: Any, instructions: Any = None) -> list[dict]: - """Responses API input -> OpenAI chat message list. - - This intentionally covers the common message/function-call shapes used by - agent SDKs. Unknown input items are preserved as user text where possible. - """ - messages: list[dict] = [] - if instructions: - messages.append({"role": "system", "content": _flatten_content(instructions)}) - - if isinstance(input_value, str): - messages.append({"role": "user", "content": input_value}) - return messages - - if not isinstance(input_value, list): - messages.append({"role": "user", "content": _flatten_content(input_value)}) - return messages - - for item in input_value: - if isinstance(item, str): - messages.append({"role": "user", "content": item}) + normalized: list[dict] = [] + for tool in tools: + if not isinstance(tool, dict): continue - if not isinstance(item, dict): - messages.append({"role": "user", "content": str(item)}) + if tool.get("type") and tool.get("type") != "function": continue - - typ = item.get("type") - if typ == "function_call_output": - messages.append( + function = tool.get("function") if isinstance(tool.get("function"), dict) else None + if function is not None: + name = function.get("name") + if not name: + continue + normalized.append( { - "role": "tool", - "tool_call_id": item.get("call_id") or item.get("id") or "", - "content": item.get("output", ""), + "type": "function", + "function": { + "name": name, + "description": function.get("description", ""), + "parameters": function.get("parameters") or {"type": "object", "properties": {}}, + }, } ) - elif typ == "function_call": - messages.append( + else: + name = tool.get("name") + if not name: + continue + normalized.append( { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": item.get("call_id") or item.get("id") or f"call_{secrets.token_hex(8)}", - "type": "function", - "function": { - "name": item.get("name", "tool"), - "arguments": item.get("arguments", "{}"), - }, - } - ], + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, + }, } ) - elif item.get("role"): - messages.append({"role": item.get("role"), "content": item.get("content", "")}) - elif typ == "message": - messages.append({"role": item.get("role", "user"), "content": item.get("content", "")}) - else: - messages.append({"role": "user", "content": _flatten_content(item)}) - return messages + return normalized or None -def _select_kind(s: Session, messages: list[dict]) -> str: - target = s.main - msg_hashes = [_hash(m) for m in messages] - if target.seen_msgs == 0: - kind = "new" - else: - is_append = len(msg_hashes) >= target.seen_msgs and msg_hashes[: target.seen_msgs] == target.msg_hashes - if is_append: - kind = "append" - else: - if target.turns: - s.segments.append(make_turn_segment(target.turns, kind="wipe")) - kind = "wipe" - return kind - - -def _replace_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: - target.chat_messages = _translate_chat_messages(messages) - target.turns.clear() - target.seen_msgs = len(messages) - target.msg_hashes = [_hash(m) for m in messages] - if tools_schema is not None: - target.tools_schema = tools_schema - - -def _extend_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: - translated = _translate_chat_messages(messages[target.seen_msgs :]) - target.chat_messages.extend(translated) - target.seen_msgs = len(messages) - target.msg_hashes = [_hash(m) for m in messages] - if tools_schema is not None: - target.tools_schema = tools_schema - - -def _build_prompt(target: Chain, messages: list[dict], tools_schema: list[dict] | None, kind: str, tok) -> list[int]: - (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, messages, tools_schema) - return render_token_ids(target, tok) - - -async def _generate( - prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None -) -> TurnRecord: - return await call_vllm_generate( - prompt_ids, - s, - body, - app, - max_token_keys=("max_output_tokens", "max_completion_tokens", "max_tokens"), - stop_keys=("stop",), - log_prefix="openai_adapter", - logger=logger, - session_id=session_id, - ) +# --- Reply building: parsed output -> OpenAI wire message + manager_message --- -def _parse_turn(target: Chain, turn: TurnRecord, app) -> ParsedModelOutput: - tok = app[TOKENIZER_KEY] - raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" - return parse_model_output( - raw_output, - tokenizer=tok, - tools_schema=target.tools_schema, - tool_parser_name=app[TOOL_PARSER_KEY], - reasoning_parser_name=app[REASONING_PARSER_KEY], - ) +def _build_reply_parts(parsed: ParsedModelOutput, finish: str) -> tuple[dict[str, Any], dict[str, Any], str]: + """Return (wire_message, manager_message, wire_finish). + wire_message follows the OpenAI Chat-Completions spec: tool_calls[].id is a + unique correlation id and tool_calls[].function.arguments is a JSON-encoded + string (clients depend on this). -def _openai_tool_calls(tool_uses: list[dict[str, Any]]) -> list[dict[str, Any]]: - calls: list[dict[str, Any]] = [] - for tool_use in tool_uses: + manager_message is the shape fed to record_turn: arguments is a dict so + chat-template replay succeeds and the manager's history match (dict equality) + holds against the echo on the next turn, and the wire-only id is omitted. + """ + wire_tool_calls: list[dict[str, Any]] = [] + manager_tool_calls: list[dict[str, Any]] = [] + for tu in parsed.tool_uses: + name = tu.get("name", "tool") + args_dict = tu.get("input") or {} + if not isinstance(args_dict, dict): + args_dict = {"_raw_arguments": str(args_dict)} call_id = f"call_{secrets.token_hex(12)}" - calls.append( + wire_tool_calls.append( { "id": call_id, "type": "function", "function": { - "name": tool_use.get("name", "tool"), - "arguments": _json_arguments(tool_use.get("input") or {}), + "name": name, + "arguments": json.dumps(args_dict, ensure_ascii=False, sort_keys=True), + }, + } + ) + manager_tool_calls.append( + { + "type": "function", + "function": { + "name": name, + "arguments": args_dict, }, } ) - return calls + wire_message: dict[str, Any] = { + "role": "assistant", + # send content=null when there are tool_calls: some OpenAI clients split + # a mixed text+tool_calls turn into two echoed messages otherwise, which + # diverges the history match against our leaf + "content": None if wire_tool_calls else (parsed.text or None), + } + # manager_message must match what the client echoes on the next request, or + # the manager's history match (dict equality) diverges and every turn forks. + # Differences from wire_message, each needed to match the echo: + # * no reasoning_content -- some clients strip it on echo (the reasoning + # token ids are still kept in the trained tokens, only the text drops) + # * only the first tool_call -- some clients drop extra parallel tool_calls + # * empty content when tool_calls are present -- mirrors content=null above + manager_message: dict[str, Any] = { + "role": "assistant", + "content": "" if wire_tool_calls else (parsed.text or ""), + } + if parsed.reasoning: + wire_message["reasoning_content"] = parsed.reasoning + if wire_tool_calls: + wire_message["tool_calls"] = wire_tool_calls[:1] + manager_message["tool_calls"] = manager_tool_calls[:1] -def _finish_reason(parsed: ParsedModelOutput, finish: str) -> str: if parsed.tool_uses: - return "tool_calls" - if finish == "length": - return "length" - return "stop" + wire_finish = "tool_calls" + elif finish == "length": + wire_finish = "length" + else: + wire_finish = "stop" + return wire_message, manager_message, wire_finish -def _chat_message(parsed: ParsedModelOutput) -> dict[str, Any]: - tool_calls = _openai_tool_calls(parsed.tool_uses) - message: dict[str, Any] = { - "role": "assistant", - "content": parsed.text if parsed.text else None, - } - if parsed.reasoning: - message["reasoning_content"] = parsed.reasoning - if tool_calls: - message["tool_calls"] = tool_calls - return message + +# --- Request framing: session id + wire response/stream rendering --- + + +def _request_session_id(request: web.Request, body: dict) -> str: + """Resolve sid: Authorization: Bearer first (where an OpenAI client + propagates its API key), then body-level hints (metadata.session_id / user).""" + return sid_from_bearer(request) or sid_from_body(body) or "default" def _usage(in_tok: int, out_tok: int) -> dict[str, int]: @@ -352,58 +298,10 @@ def _usage(in_tok: int, out_tok: int) -> dict[str, int]: } -def _responses_usage(in_tok: int, out_tok: int) -> dict[str, int]: - return { - "input_tokens": in_tok, - "output_tokens": out_tok, - "total_tokens": in_tok + out_tok, - } - - -def _request_session_id(request: web.Request, body: dict) -> str: - return request_session_id(request, body=body) - - -async def _run_turn( - request: web.Request, body: dict, messages: list[dict] -) -> tuple[TurnRecord, ParsedModelOutput, int, int]: - sid = _request_session_id(request, body) - adapter = request.app[ADAPTER_KEY] - if sid in adapter.closed: - raise web.HTTPServiceUnavailable(text="session closed") - app = request.app - s = adapter.store.setdefault(sid, Session()) - task = asyncio.current_task() - adapter.inflight.setdefault(sid, set()).add(task) - try: - async with s.lock: - target = s.main - tools_schema = _normalize_tools(body.get("tools")) - kind = _select_kind(s, messages) - prompt_ids = _build_prompt(target, messages, tools_schema, kind, app[TOKENIZER_KEY]) - turn = await _generate(prompt_ids, s, body, app, session_id=sid) - parsed = _parse_turn(target, turn, app) - target.turns.append(turn) - return turn, parsed, len(prompt_ids), len(turn.output_ids) - finally: - adapter.inflight.get(sid, set()).discard(task) - - -async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: - body = await request.json() - messages = body.get("messages") or [] - if not isinstance(messages, list): - raise web.HTTPBadRequest(text="messages must be a list") - turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) - if body.get("stream"): - return await _stream_chat_completion(request, body, parsed, turn.finish_reason, in_tok, out_tok) - return web.json_response(_chat_completion_response(body, parsed, turn.finish_reason, in_tok, out_tok)) - - -def _chat_completion_response( +def _render_response( body: dict, - parsed: ParsedModelOutput, - finish: str, + wire_message: dict[str, Any], + wire_finish: str, in_tok: int, out_tok: int, ) -> dict[str, Any]: @@ -415,22 +313,29 @@ def _chat_completion_response( "choices": [ { "index": 0, - "message": _chat_message(parsed), - "finish_reason": _finish_reason(parsed, finish), + "message": wire_message, + "finish_reason": wire_finish, } ], "usage": _usage(in_tok, out_tok), } -async def _stream_chat_completion( +async def _render_stream( request: web.Request, body: dict, - parsed: ParsedModelOutput, - finish: str, + wire_message: dict[str, Any], + wire_finish: str, in_tok: int, out_tok: int, ) -> web.StreamResponse: + """Emit the OpenAI Chat-Completions SSE stream. + + Each chunk has the shape `data: {chatcmpl ...}\n\n`, ending with + `data: [DONE]`. The whole turn is realised on the server before streaming + (we have no token-level deltas from vllm here), so we emit one role chunk, + then content / reasoning / tool_calls in single delta chunks each. + """ out = web.StreamResponse( status=200, headers={ @@ -443,131 +348,31 @@ async def _stream_chat_completion( completion_id = f"chatcmpl_{secrets.token_hex(12)}" created = int(time.time()) - async def emit(choice_delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: + async def emit(delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": body.get("model", "vime-actor"), - "choices": [{"index": 0, "delta": choice_delta, "finish_reason": finish_reason}], + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], } if usage is not None: chunk["usage"] = usage await out.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()) await emit({"role": "assistant"}) - if parsed.reasoning: - await emit({"reasoning_content": parsed.reasoning}) - if parsed.text: - await emit({"content": parsed.text}) - for idx, call in enumerate(_openai_tool_calls(parsed.tool_uses)): - await emit({"tool_calls": [{**call, "index": idx}]}) - await emit({}, finish_reason=_finish_reason(parsed, finish), usage=_usage(in_tok, out_tok)) + reasoning = wire_message.get("reasoning_content") + if reasoning: + await emit({"reasoning_content": reasoning}) + content = wire_message.get("content") + if content: + await emit({"content": content}) + # emit all tool_calls in a single chunk: some clients accumulate per-index + # arguments fragments across chunks, collapsing N parallel tool_calls into + # one call with a concatenated (and unparseable) arguments string + tool_calls = wire_message.get("tool_calls") or [] + if tool_calls: + await emit({"tool_calls": [{**call, "index": idx} for idx, call in enumerate(tool_calls)]}) + await emit({}, finish_reason=wire_finish, usage=_usage(in_tok, out_tok)) await out.write(b"data: [DONE]\n\n") return out - - -async def _handle_responses(request: web.Request) -> web.StreamResponse: - body = await request.json() - messages = _responses_input_to_messages(body.get("input", ""), body.get("instructions")) - turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) - if body.get("stream"): - return await _stream_response(request, body, parsed, turn.finish_reason, in_tok, out_tok) - return web.json_response(_response_response(body, parsed, turn.finish_reason, in_tok, out_tok)) - - -def _response_output(parsed: ParsedModelOutput) -> list[dict[str, Any]]: - output: list[dict[str, Any]] = [] - if parsed.reasoning: - output.append( - { - "id": f"rs_{secrets.token_hex(12)}", - "type": "reasoning", - "summary": [{"type": "summary_text", "text": parsed.reasoning}], - } - ) - if parsed.text: - output.append( - { - "id": f"msg_{secrets.token_hex(12)}", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": parsed.text, "annotations": []}], - } - ) - for call in _openai_tool_calls(parsed.tool_uses): - output.append( - { - "id": f"fc_{secrets.token_hex(12)}", - "type": "function_call", - "status": "completed", - "call_id": call["id"], - "name": call["function"]["name"], - "arguments": call["function"]["arguments"], - } - ) - if not output: - output.append( - { - "id": f"msg_{secrets.token_hex(12)}", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "", "annotations": []}], - } - ) - return output - - -def _response_response( - body: dict, - parsed: ParsedModelOutput, - finish: str, - in_tok: int, - out_tok: int, -) -> dict[str, Any]: - status = "incomplete" if finish == "length" else "completed" - return { - "id": f"resp_{secrets.token_hex(12)}", - "object": "response", - "created_at": int(time.time()), - "status": status, - "model": body.get("model", "vime-actor"), - "output": _response_output(parsed), - "usage": _responses_usage(in_tok, out_tok), - } - - -async def _stream_response( - request: web.Request, - body: dict, - parsed: ParsedModelOutput, - finish: str, - in_tok: int, - out_tok: int, -) -> web.StreamResponse: - out = web.StreamResponse( - status=200, - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Connection": "keep-alive", - }, - ) - await out.prepare(request) - response = _response_response(body, parsed, finish, in_tok, out_tok) - created = {"type": "response.created", "response": response} - await out.write(f"event: response.created\ndata: {json.dumps(created, ensure_ascii=False)}\n\n".encode()) - if parsed.text: - delta = {"type": "response.output_text.delta", "delta": parsed.text} - await out.write( - f"event: response.output_text.delta\ndata: {json.dumps(delta, ensure_ascii=False)}\n\n".encode() - ) - completed = {"type": "response.completed", "response": response} - await out.write(f"event: response.completed\ndata: {json.dumps(completed, ensure_ascii=False)}\n\n".encode()) - return out - - -async def _ok(request: web.Request) -> web.Response: - return await ok_response(request) diff --git a/examples/coding_agent_rl/aiohttp_threaded.py b/vime/agent/aiohttp_threaded.py similarity index 87% rename from examples/coding_agent_rl/aiohttp_threaded.py rename to vime/agent/aiohttp_threaded.py index a5a17652d..ad5fa0ccb 100644 --- a/examples/coding_agent_rl/aiohttp_threaded.py +++ b/vime/agent/aiohttp_threaded.py @@ -8,6 +8,18 @@ from typing import Any from aiohttp import web +from aiohttp.web_log import AccessLogger + + +class FilteredAccessLogger(AccessLogger): + SLOW_THRESHOLD_SEC = 120.0 + + def log(self, request, response, time): + if request.method == "HEAD": + return + if response.status == 200 and time <= self.SLOW_THRESHOLD_SEC: + return + super().log(request, response, time) @dataclass @@ -18,10 +30,6 @@ class AppHandle: loop: asyncio.AbstractEventLoop runner: web.AppRunner - @property - def url(self) -> str: - return f"http://{self.host}:{self.port}" - def stop(self) -> None: async def _shutdown() -> None: await self.runner.cleanup() diff --git a/vime/agent/harness/__init__.py b/vime/agent/harness/__init__.py new file mode 100644 index 000000000..caac42b06 --- /dev/null +++ b/vime/agent/harness/__init__.py @@ -0,0 +1,14 @@ +"""Swappable coding-agent harnesses (Claude Code, Codex, ...).""" + +from __future__ import annotations + +from .claude_code import ClaudeCodeHarness +from .codex import CodexHarness +from .common import BaseHarness, HarnessContext + +__all__ = [ + "BaseHarness", + "HarnessContext", + "ClaudeCodeHarness", + "CodexHarness", +] diff --git a/vime/agent/harness/claude_code.py b/vime/agent/harness/claude_code.py new file mode 100644 index 000000000..11f0ad3fc --- /dev/null +++ b/vime/agent/harness/claude_code.py @@ -0,0 +1,71 @@ +"""Claude Code harness.""" + +from __future__ import annotations + +import json +import os +import shlex +from pathlib import Path + +from vime.agent.sandbox import Sandbox + +from .common import BaseHarness, HarnessContext, install_npm_cli, run_agent + + +class ClaudeCodeHarness(BaseHarness): + name = "claude_code" + + # host paths + CLI knobs, all under the agent-layer VIME_AGENT_* prefix + node_tarball_env = "VIME_AGENT_NODE_TARBALL" + cli_tarball_env = "VIME_AGENT_CC_TARBALL" + extra_args_env = "VIME_AGENT_CC_EXTRA_ARGS" + extra_envs_env = "VIME_AGENT_CC_EXTRA_ENVS" + + launch_flags = ( + "--permission-mode bypassPermissions " + "--output-format stream-json --include-partial-messages " + "--include-hook-events --verbose" + ) + + static_env = { + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + } + + async def install_cli(self, sb: Sandbox) -> None: + await install_npm_cli( + sb, + node_runtime=Path(os.environ[self.node_tarball_env]), + npm_package=Path(os.environ[self.cli_tarball_env]), + check_cmd="ls -la /usr/local/bin/claude && /usr/local/bin/claude --version", + ) + + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + """Pre-ack bypass-permissions so claude-code starts headless.""" + settings = json.dumps({"hasCompletedOnboarding": True, "bypassPermissionsModeAccepted": True}) + await sb.exec( + "mkdir -p /home/agent/.claude && " + f"echo {shlex.quote(settings)} " + "| tee /home/agent/.claude.json /home/agent/.claude/settings.json > /dev/null && " + "chown -R agent:agent /home/agent/.claude /home/agent/.claude.json", + user="root", + check=True, + timeout=60, + ) + + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + cmd = f"/usr/local/bin/claude -p {shlex.quote(prompt)} {self.launch_flags}" + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + cmd = f"{cmd} {extra}" + env = { + "ANTHROPIC_BASE_URL": ctx.adapter_url, + "ANTHROPIC_AUTH_TOKEN": ctx.session_id, + "ANTHROPIC_MODEL": ctx.model_label, + **self.static_env, + } + extra_envs = os.environ.get(self.extra_envs_env, "").strip() + if extra_envs: + env.update(json.loads(extra_envs)) + return await run_agent(sb, workdir=ctx.workdir, start_cmd=cmd, env=env, time_budget_sec=time_budget_sec) diff --git a/vime/agent/harness/codex.py b/vime/agent/harness/codex.py new file mode 100644 index 000000000..a913e19e4 --- /dev/null +++ b/vime/agent/harness/codex.py @@ -0,0 +1,86 @@ +"""Codex harness. + +Two non-obvious bits: the provider base_url must be inline in the TOML (Codex +only honours env vars for the default OpenAI provider), and the config is written +via a base64 round-trip to dodge shell-quoting traps. +""" + +from __future__ import annotations + +import base64 +import json +import os +import shlex +from pathlib import Path + +from vime.agent.sandbox import Sandbox + +from .common import BaseHarness, HarnessContext, install_npm_cli, run_agent + + +class CodexHarness(BaseHarness): + name = "codex" + + # host paths + CLI knobs, all under the agent-layer VIME_AGENT_* prefix + node_tarball_env = "VIME_AGENT_NODE_TARBALL" + cli_tarball_env = "VIME_AGENT_CODEX_TARBALL" + extra_args_env = "VIME_AGENT_CODEX_EXTRA_ARGS" + extra_envs_env = "VIME_AGENT_CODEX_EXTRA_ENVS" + + # static flags after ``codex exec``; --skip-git-repo-check lets it run in + # workdirs whose git check is brittle (e.g. shallow clones) + exec_flags = "--skip-git-repo-check" + + # config.toml written into the sandbox. base_url MUST be inline here (Codex + # only honours env vars for the default OpenAI provider). {model} / {base_url} + # are filled per run in write_config; the rest is fixed wiring. + config_toml = ( + 'model = "{model}"\n' + 'model_provider = "vime"\n' + "\n" + "[model_providers.vime]\n" + 'name = "vime"\n' + 'base_url = "{base_url}"\n' + 'env_key = "OPENAI_API_KEY"\n' + 'wire_api = "chat"\n' + ) + + async def install_cli(self, sb: Sandbox) -> None: + await install_npm_cli( + sb, + node_runtime=Path(os.environ[self.node_tarball_env]), + npm_package=Path(os.environ[self.cli_tarball_env]), + check_cmd="codex --version", + ) + + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + toml = self.config_toml.format(model=ctx.model_label, base_url=f"{ctx.adapter_url}/v1") + toml_b64 = base64.b64encode(toml.encode("utf-8")).decode("ascii") + await sb.exec( + "mkdir -p /home/agent/.codex && " + # base64 round-trip avoids any single-quote / heredoc shell-quoting trap + f"echo {shlex.quote(toml_b64)} | base64 -d > /home/agent/.codex/config.toml && " + "chown -R agent:agent /home/agent/.codex", + user="root", + check=True, + timeout=60, + ) + + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + # ``codex exec`` is the non-interactive entrypoint + cmd = f"codex exec {self.exec_flags} {shlex.quote(prompt)}" + extra = os.environ.get(self.extra_args_env, "").strip() + if extra: + cmd = f"{cmd} {extra}" + env = { + # Codex propagates OPENAI_API_KEY into Authorization: Bearer; the + # vime adapter resolves the sid from that header. + "OPENAI_API_KEY": ctx.session_id, + "OPENAI_BASE_URL": f"{ctx.adapter_url}/v1", + } + # extra env vars as a JSON object, merged last so callers can override + # the defaults above + extra_envs = os.environ.get(self.extra_envs_env, "").strip() + if extra_envs: + env.update(json.loads(extra_envs)) + return await run_agent(sb, workdir=ctx.workdir, start_cmd=cmd, env=env, time_budget_sec=time_budget_sec) diff --git a/vime/agent/harness/common.py b/vime/agent/harness/common.py new file mode 100644 index 000000000..ca337155a --- /dev/null +++ b/vime/agent/harness/common.py @@ -0,0 +1,178 @@ +"""Harness-agnostic coding-agent lifecycle in a sandbox. + +A harness is a swappable coding agent (Claude Code, Codex, ...). Each one +installs a CLI, writes its own config, and runs the agent against a prompt. The +shared parts (create the agent user, the run skeleton, the launch-detached-and- +poll transport) live here; adding a CLI-style harness means subclassing +BaseHarness and implementing install_cli, write_config and launch_and_wait. +Two module-level helpers cover the common cases: install_npm_cli for +npm-packaged CLIs, and run_agent for the launch-the-agent-to-completion case. + +The base knows nothing about the task: run() takes only generic fields +(workdir / session_id / adapter_url / prompt). Task-specific workspace prep and +scoring live in the example layer. +""" + +from __future__ import annotations + +import asyncio +import lzma +import os +import shutil +import tempfile +from abc import ABC, ABCMeta, abstractmethod +from dataclasses import dataclass +from pathlib import Path + +from vime.agent import sandbox as _sandbox +from vime.agent.sandbox import Sandbox, exec_and_wait +from vime.utils.misc import SingletonMeta + + +class SingletonABCMeta(ABCMeta, SingletonMeta): + pass + + +# In-sandbox retry budget for the npm global install (transient flakes like +# exit 217). Cheaper than a full sandbox recreate by the caller. +NPM_INSTALL_RETRIES = 3 +NPM_INSTALL_BACKOFF_SEC = 2.0 + + +@dataclass(frozen=True) +class HarnessContext: + """Generic run context, free of any task fields. + + model_label is the model name the harness advertises to its CLI. The vime + adapter ignores it and serves whatever upstream vllm has loaded, so it is + not a run() parameter. + """ + + workdir: str + session_id: str + adapter_url: str + model_label: str = "vime-actor" + + +class BaseHarness(ABC, metaclass=SingletonABCMeta): + """Base lifecycle for a sandbox-resident coding agent.""" + + # short identifier set by each subclass (claude_code / codex) + name: str = "" + + @abstractmethod + async def install_cli(self, sb: Sandbox) -> None: + """Install the harness CLI into the sandbox. + npm-packaged harnesses delegate to install_npm_cli.""" + + @abstractmethod + async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None: + """Write any CLI config files into the sandbox.""" + + @abstractmethod + async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int: + """Run the agent to completion and return its exit code. + + A non-interactive CLI builds one shell command and hands it to + run_agent. An interactive or long-running harness drives its own loop + here instead. + """ + + async def run( + self, + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + time_budget_sec: int, + prompt: str, + ) -> int: + """Run the harness in the sandbox and return its exit code. + + Steps: ensure the agent user -> write config -> launch and wait. + Workspace prep (writing the problem statement etc.) is the caller's job + and must run before this. + """ + await _sandbox.ensure_agent_user(sb, workdir) + ctx = HarnessContext( + workdir=workdir, + session_id=session_id, + adapter_url=adapter_url, + ) + await self.write_config(sb, ctx) + return await self.launch_and_wait(sb, ctx, prompt, time_budget_sec) + + +async def run_agent(sb: Sandbox, *, workdir: str, start_cmd: str, env: dict[str, str], time_budget_sec: int) -> int: + """Launch the agent (start_cmd) and run it to completion, returning its exit code.""" + meta_dir = f"{workdir}/.harness" + await sb.exec(f"mkdir -p {meta_dir} && chown agent:agent {meta_dir}", user="root", check=True, timeout=30) + exit_code, _ = await exec_and_wait( + sb, + cmd=start_cmd, + user="agent", + env=env, + workdir=workdir, + out_file=f"{meta_dir}/trajectory.jsonl", + time_budget_sec=time_budget_sec, + tag="run", + want_output=False, + ) + return exit_code + + +async def install_npm_cli( + sb: Sandbox, + *, + node_runtime: Path, + npm_package: Path, + check_cmd: str, +) -> None: + """Install an npm-packaged CLI into the sandbox: the Node 22 runtime first, + then the CLI's npm package (global install, then self-check via check_cmd). + Non-npm harnesses write their own install_cli.""" + await install_node22(sb, node_runtime) + + await sb.write_file("/tmp/harness-cli.tgz", npm_package) + install_cmd = "npm install -g --prefix=/usr/local --no-audit --no-fund /tmp/harness-cli.tgz && " + check_cmd + # Detached install with a few in-place retries for transient disk flakes. + last_log = "" + for attempt in range(NPM_INSTALL_RETRIES): + exit_code, last_log = await exec_and_wait( + sb, cmd=install_cmd, user="root", time_budget_sec=300, tag="harness-npm-install" + ) + if exit_code == 0: + return + if attempt + 1 < NPM_INSTALL_RETRIES: + await asyncio.sleep(NPM_INSTALL_BACKOFF_SEC * (attempt + 1)) + raise RuntimeError( + f"npm install failed after {NPM_INSTALL_RETRIES} attempts (exit={exit_code}):\n{last_log[-1000:]}" + ) + + +async def install_node22(sb: Sandbox, host_tarball: Path) -> None: + """Install Node 22 over the base image (some base images ship a version too + old for the CLI). A .xz tarball is decompressed on the host (cached) so + sandboxes without xz-utils can still run a plain `tar xf`.""" + host_tarball = Path(host_tarball) + if host_tarball.suffix == ".xz": + plain = Path(tempfile.gettempdir()) / f"coding_agent_rl.{host_tarball.stem}.tar" + if not plain.exists(): + tmp = plain.with_suffix(".tar.partial") + with lzma.open(host_tarball, "rb") as src, open(tmp, "wb") as dst: + shutil.copyfileobj(src, dst) + os.replace(tmp, plain) + host_tarball = plain + await sb.write_file("/tmp/node22.tar", host_tarball) + await sb.exec( + "set -e && mkdir -p /opt/node22 && " + "tar xf /tmp/node22.tar -C /opt/node22 --strip-components=1 && " + "ln -sf /opt/node22/bin/node /usr/local/bin/node && " + "ln -sf /opt/node22/bin/npm /usr/local/bin/npm && " + "ln -sf /opt/node22/bin/npx /usr/local/bin/npx && " + "hash -r 2>/dev/null || true && node --version && npm --version", + user="root", + timeout=180, + check=True, + ) diff --git a/vime/agent/parsing.py b/vime/agent/parsing.py index 86d48a26d..7df615dec 100644 --- a/vime/agent/parsing.py +++ b/vime/agent/parsing.py @@ -19,12 +19,13 @@ class ParsedModelOutput: reasoning: str text: str tool_uses: list[dict[str, Any]] + ill_formed: bool = False def parse_model_output( raw_output: str, *, - tokenizer, + tokenizer=None, tools_schema: list[dict] | None, tool_parser_name: str | None, reasoning_parser_name: str | None, @@ -46,11 +47,12 @@ def parse_model_output( if not reasoning and "" in body_text: reasoning, body_text = body_text.split("", 1) - body_text, tool_uses = parse_tool_uses(body_text, tools_schema, tool_parser_name, tokenizer) + body_text, tool_uses, ill_formed = parse_tool_uses(body_text, tools_schema, tool_parser_name, tokenizer) return ParsedModelOutput( reasoning=reasoning, text=(body_text or "").strip(), tool_uses=tool_uses, + ill_formed=ill_formed, ) @@ -59,9 +61,10 @@ def parse_tool_uses( tools_schema: list[dict] | None, tool_parser_name: str | None, tokenizer, -) -> tuple[str, list[dict[str, Any]]]: +) -> tuple[str, list[dict[str, Any]], bool]: """Parse tool calls from body text and return visible text plus tool uses.""" tool_uses: list[dict[str, Any]] = [] + ill_formed = False if tool_parser_name and tools_schema: from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tool_parsers import ToolParserManager @@ -80,12 +83,13 @@ def parse_tool_uses( args = json.loads(call.function.arguments or "{}") except json.JSONDecodeError: args = {"_raw_arguments": call.function.arguments} + ill_formed = True tool_uses.append({"name": call.function.name or "tool", "input": args}) if not tool_uses and tools_schema: body_text, tool_uses = parse_xml_tool_uses(body_text, tools_schema) - return body_text, tool_uses + return body_text, tool_uses, ill_formed def parse_xml_tool_uses(body_text: str, tools_schema: list[dict]) -> tuple[str, list[dict[str, Any]]]: diff --git a/vime/agent/sandbox.py b/vime/agent/sandbox.py index 6447bd2ee..6ab7bf7f5 100644 --- a/vime/agent/sandbox.py +++ b/vime/agent/sandbox.py @@ -10,9 +10,10 @@ import asyncio import io -import json import logging import os +import random +import time from pathlib import Path from typing import Protocol, runtime_checkable @@ -29,6 +30,11 @@ class Sandbox(Protocol): ``write_file`` accepts either in-memory content (``str``/``bytes``) or a host ``Path`` to stream into the sandbox. + + ``idempotent`` is a hint for the backend's transport-retry policy: callers + mark whether re-sending the command after a severed response is safe to + replay (see ``E2BSandbox._rpc_retry``). Backends without retries may + ignore it. """ sandbox_id: str @@ -45,6 +51,7 @@ async def exec( env: dict[str, str] | None = None, timeout: int = 120, check: bool = False, + idempotent: bool = True, ) -> ExecResult: ... async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: ... @@ -52,7 +59,97 @@ async def write_file(self, sandbox_path: str, content: FileContent, *, user: str async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: ... +EXIT_TIME_BUDGET_EXCEEDED = -1 + + +async def _await_done_marker(sb: Sandbox, done_file: str, *, user: str, time_budget_sec: int) -> int: + """Poll a detached command's exit-code marker until it appears, returning the + exit code (or ``EXIT_TIME_BUDGET_EXCEEDED`` if the budget runs out first). + + The 5s ``test -f && cat`` polls are deliberately short, idempotent RPCs -- + they keep the sandbox alive against idle GC while the detached command runs + over a stream the gateway can't sever. + """ + deadline = time.time() + time_budget_sec + while time.time() < deadline: + await asyncio.sleep(5) + ec, out, _ = await sb.exec(f"test -f {done_file} && cat {done_file}", user=user, timeout=15, check=False) + if ec == 0 and (out or "").strip(): + return int(out.strip()) + return EXIT_TIME_BUDGET_EXCEEDED + + +async def exec_and_wait( + sb: Sandbox, + *, + cmd: str, + time_budget_sec: int, + tag: str, + user: str = "root", + env: dict[str, str] | None = None, + workdir: str | None = None, + out_file: str | None = None, + want_output: bool = False, +) -> tuple[int, str]: + """Run ``cmd`` to completion detached, returning ``(exit_code, output)``. + + A plain ``sb.exec`` keeps an HTTP/2 stream open for the command's whole + runtime, so a long-running command (build, test suite) outlives what the + E2B gateway will hold a single response stream open for: the stream gets + severed mid-run and we lose the exit code with no safe way to retry a + non-idempotent command. Instead we ``setsid`` the command fully detached, + redirect its output to a file, and have it drop its exit code into a marker + file. The caller side then becomes a sequence of short, idempotent RPCs -- + write the launcher, fire-and-forget the spawn, then poll for the marker (see + ``_await_done_marker``) -- none of which depend on a stream staying alive, + and the polling doubles as an idle-GC keepalive while the command runs. + """ + out_file = out_file or f"/tmp/.{tag}.out" + done_file = f"/tmp/.{tag}.done" + launcher = f"/tmp/.{tag}.sh" + lock_dir = f"/tmp/.{tag}.spawned" + prefix = f"cd {workdir}\nexport HOME=/home/{user}\n" if workdir else "" + launcher_body = f"#!/bin/bash\n{prefix}{cmd}\necho $? > {done_file}\n" + await sb.write_file(launcher, launcher_body, user=user) + + # Clear the previous invocation's state in its own idempotent RPC, *before* + # the guarded spawn. The mkdir guard below exists only to dedupe transport + # retries of this one spawn (a severed response replayed by _rpc_retry); it + # must not survive into the next logical invocation of the same tag (e.g. + # install_npm_cli's retry loop), which would skip the spawn entirely and + # read the previous run's stale exit-code marker. Callers must not overlap + # two exec_and_wait calls with the same tag. + await sb.exec( + f"rm -rf {lock_dir}; rm -f {out_file} {done_file}", + user=user, + timeout=30, + check=True, + idempotent=True, + ) + await sb.exec( + f"chmod +x {launcher}; " + f"mkdir {lock_dir} 2>/dev/null || exit 0; " + f"setsid bash {launcher} < /dev/null > {out_file} 2>&1 &", + user=user, + env=env, + timeout=30, + check=True, + idempotent=True, + ) + exit_code = await _await_done_marker(sb, done_file, user=user, time_budget_sec=time_budget_sec) + if exit_code == 0 and not want_output: + return exit_code, "" + if want_output: + return exit_code, await sb.read_file(out_file, user=user) + _, tail, _ = await sb.exec(f"tail -c 512 {out_file} 2>/dev/null", user=user, timeout=15, check=False) + return exit_code, tail or "" + + def _getenv(*names: str, default: str = "") -> str: + """First non-empty environment value among ``names`` (else ``default``). + + Lets a setting carry a primary name plus legacy aliases: list the canonical + ``VIME_AGENT_*`` name first, older names after.""" for name in names: value = os.environ.get(name) if value is not None and value.strip(): @@ -63,60 +160,34 @@ def _getenv(*names: str, default: str = "") -> str: class E2BSandbox: """Async context manager around e2b.AsyncSandbox.""" - metadata_file_env = ("VIME_AGENT_SANDBOX_METADATA_FILE", "SWE_SANDBOX_METADATA_FILE") - metadata_json_env = ("VIME_AGENT_SANDBOX_METADATA_JSON", "SWE_SANDBOX_METADATA_JSON") image_metadata_key_env = ("VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY", "SWE_SANDBOX_IMAGE_METADATA_KEY") lifetime_sec_env = ("VIME_AGENT_SANDBOX_LIFETIME_SEC", "SWE_SANDBOX_LIFETIME_SEC") rpc_retries_env = ("VIME_AGENT_SANDBOX_RPC_RETRIES", "SWE_RPC_RETRIES") + size_env = ("VIME_AGENT_E2B_SANDBOX_SIZE", "SWE_E2B_SANDBOX_SIZE") default_lifetime_sec = 3600 - default_rpc_retries = 3 - # With retries=3 the sleep budget is 3s, which handles common E2B h2 reset - # / SSL / pool-timeout flaps without stalling rollout steps for too long. + default_rpc_retries = 6 + default_size = "md" rpc_backoff_base_sec = 1.0 + rpc_backoff_cap_sec = 32.0 def __init__( self, image: str, *, timeout: int | None = None, - metadata: dict[str, str] | None = None, image_metadata_key: str | None = None, rpc_retries: int | None = None, + size: str | None = None, ) -> None: self.image = image self.timeout = timeout if timeout is not None else self._lifetime_sec_from_env() - self.metadata = dict(metadata) if metadata is not None else self._metadata_from_env() self.image_metadata_key = image_metadata_key or self._image_metadata_key_from_env() self.rpc_retries = rpc_retries if rpc_retries is not None else self._rpc_retries_from_env() + self.size = size if size is not None else self._size_from_env() self._sb = None self.sandbox_id = "" - @classmethod - def _metadata_from_env(cls) -> dict[str, str]: - """Read E2B routing metadata from file or JSON environment values.""" - file_path = _getenv(*cls.metadata_file_env) - raw = "" - if file_path: - try: - raw = Path(file_path).read_text() - except OSError as e: - logger.warning("[agent.sandbox] metadata file %s unreadable: %s", file_path, e) - raw = "" - if not raw: - raw = _getenv(*cls.metadata_json_env) - if not raw: - return {} - try: - md = json.loads(raw) - except json.JSONDecodeError as e: - logger.warning("[agent.sandbox] metadata not valid JSON, ignoring: %s", e) - return {} - if not isinstance(md, dict): - logger.warning("[agent.sandbox] metadata must be a JSON object, got %s", type(md).__name__) - return {} - return {str(k): str(v) for k, v in md.items()} - @classmethod def _image_metadata_key_from_env(cls) -> str | None: return _getenv(*cls.image_metadata_key_env) or None @@ -129,11 +200,13 @@ def _lifetime_sec_from_env(cls) -> int: def _rpc_retries_from_env(cls) -> int: return int(_getenv(*cls.rpc_retries_env, default=str(cls.default_rpc_retries))) - @staticmethod - def _is_transient_rpc_error(e: BaseException) -> bool: - """True if e is a transient E2B client-side failure safe to retry.""" - name = type(e).__name__ - if name in { + @classmethod + def _size_from_env(cls) -> str: + return _getenv(*cls.size_env, default=cls.default_size) + + # Transient client-side failures safe to retry. + _TRANSIENT_RPC_ERRORS = frozenset( + { "ProtocolError", "LocalProtocolError", "WriteError", @@ -145,7 +218,14 @@ def _is_transient_rpc_error(e: BaseException) -> bool: "PoolTimeout", "RemoteProtocolError", "SSLError", - }: + } + ) + + @classmethod + def _is_transient_rpc_error(cls, e: BaseException) -> bool: + """True if e is a transient E2B client-side failure safe to retry.""" + name = type(e).__name__ + if name in cls._TRANSIENT_RPC_ERRORS: return True msg = str(e) if name == "SandboxException": @@ -154,8 +234,15 @@ def _is_transient_rpc_error(e: BaseException) -> bool: return True return False - async def _rpc_retry(self, op_name: str, coro_factory): - """Run coro_factory() with retries for transient E2B RPC failures.""" + async def _rpc_retry(self, op_name: str, coro_factory, *, idempotent: bool = True): + """Run coro_factory() with retries for transient E2B RPC failures. + + :param idempotent: When False, a transient failure is re-raised instead + of retried: re-running a non-idempotent op (e.g. a process-spawning + exec) after a severed response could double-execute it. Idempotent + ops (the default: create / read_file / write_file / short read-only + execs) retry as before. + """ last_err = None for attempt in range(self.rpc_retries): try: @@ -163,9 +250,13 @@ async def _rpc_retry(self, op_name: str, coro_factory): except Exception as e: if not self._is_transient_rpc_error(e): raise + if not idempotent: + raise last_err = e if attempt + 1 < self.rpc_retries: - backoff = self.rpc_backoff_base_sec * (2**attempt) + await self._reset_conn_pool() + ceiling = min(self.rpc_backoff_cap_sec, self.rpc_backoff_base_sec * (2**attempt)) + backoff = random.uniform(0.0, ceiling) logger.debug( "[agent.sandbox] %s transient %s, retry %d/%d in %.1fs: %s", op_name, @@ -179,6 +270,14 @@ async def _rpc_retry(self, op_name: str, coro_factory): assert last_err is not None raise last_err + async def _reset_conn_pool(self) -> None: + """Tear down the sandbox's httpcore pool so the next RPC reconnects.""" + try: + pool = self._sb._transport.pool # httpcore.AsyncConnectionPool + await pool.aclose() + except Exception as e: + logger.debug("[agent.sandbox] conn-pool reset skipped: %s", e) + async def __aenter__(self) -> E2BSandbox: if self.image_metadata_key is None: raise RuntimeError( @@ -189,9 +288,14 @@ async def __aenter__(self) -> E2BSandbox: ) from e2b import AsyncSandbox # type: ignore - md = dict(self.metadata) - md.setdefault(self.image_metadata_key, self.image) - self._sb = await AsyncSandbox.create(timeout=self.timeout, metadata=md) + md = {self.image_metadata_key: self.image} + + if self.size: + prefix = self.image_metadata_key.rsplit("/", 1)[0] if "/" in self.image_metadata_key else "" + size_key = f"{prefix}/size" if prefix else "size" + md[size_key] = self.size + + self._sb = await self._rpc_retry("create", lambda: AsyncSandbox.create(timeout=self.timeout, metadata=md)) self.sandbox_id = self._sb.sandbox_id return self @@ -210,6 +314,7 @@ async def exec( env: dict[str, str] | None = None, timeout: int = 120, check: bool = False, + idempotent: bool = True, ) -> ExecResult: from e2b.sandbox.commands.command_handle import CommandExitException @@ -224,6 +329,7 @@ async def exec( on_stdout=lambda s: None, on_stderr=lambda s: None, ), + idempotent=idempotent, ) return res.exit_code, res.stdout or "", res.stderr or "" except CommandExitException as e: @@ -279,3 +385,15 @@ async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: ) except Exception: return "" + + +async def ensure_agent_user(sb: Sandbox, workdir: str) -> None: + """Create the unprivileged 'agent' user that owns workdir + can git diff.""" + await sb.exec( + f"id agent >/dev/null 2>&1 || useradd -m -s /bin/bash agent && " + f"chown -R agent:agent /home/agent {workdir} && " + f"git config --system --add safe.directory '*' && id agent", + user="root", + check=True, + timeout=60, + ) diff --git a/vime/agent/trajectory.py b/vime/agent/trajectory.py index b30b69080..9181f3cb7 100644 --- a/vime/agent/trajectory.py +++ b/vime/agent/trajectory.py @@ -1,208 +1,508 @@ -"""Token-level trajectory helpers for agent rollouts.""" +"""Build a per-session training trajectory from multi-turn conversation data. + +The :class:`TrajectoryManager` builds one trajectory per session. ``record_turn`` +feeds in each turn (prompt messages + the served model's vllm snapshot), +routing it into a per-sid message tree; ``get_trajectory`` then linearizes that +tree into a ``list[Sample]`` of loss-masked training rows, tolerating TITO +re-tokenization drift via fork/replace. +""" from __future__ import annotations -import copy import dataclasses +import enum import logging +from collections.abc import Iterator from typing import Any from vime.utils.types import Sample - logger = logging.getLogger(__name__) +# =========================================================================== +# TurnRecord +# =========================================================================== + + @dataclasses.dataclass(frozen=True) class TurnRecord: - """Exact token snapshot for one assistant generation. - - ``prompt_ids`` is the full tokenized prompt sent to the generator for that - turn. ``output_ids`` is the raw generated output, and - ``output_log_probs`` is aligned with it when the rollout engine returns - per-token log probabilities. - """ + """One vllm ``/inference/v1/generate`` snapshot: the contract between an adapter and the + manager. Adapters build it from a turn's prompt/output token ids; ``record_turn`` + consumes it.""" prompt_ids: list[int] output_ids: list[int] finish_reason: str output_log_probs: list[float] = dataclasses.field(default_factory=list) + ill_formed: bool = False -@dataclasses.dataclass(frozen=True) -class TokenSegment: - """One training segment assembled from an agent trajectory.""" +# =========================================================================== +# MessageNode +# =========================================================================== - prompt_ids: list[int] - response_ids: list[int] - loss_mask: list[int] - rollout_log_probs: list[float] = dataclasses.field(default_factory=list) - metadata: dict[str, Any] = dataclasses.field(default_factory=dict) +class MessageNode: + """One node in a session's routing tree, carrying a single chat message + (``None`` for the dummy root and for an assistant leaf we generated but + whose ``response_message`` was empty). -@dataclasses.dataclass(frozen=True) -class TurnSegment: - """A frozen group of turns before token-level merge.""" - - turns: list[TurnRecord] - metadata: dict[str, Any] = dataclasses.field(default_factory=dict) - - -def make_turn_segment( - turns: list[TurnRecord], - *, - kind: str = "", - metadata: dict[str, Any] | None = None, -) -> TurnSegment: - """Freeze turns and attach conventional segment metadata.""" - frozen_turns = list(turns) - segment_metadata = dict(metadata or {}) - if kind: - segment_metadata.setdefault("segment_kind", kind) - segment_metadata.setdefault("finish_reason", frozen_turns[-1].finish_reason if frozen_turns else "") - return TurnSegment(turns=frozen_turns, metadata=segment_metadata) - - -def _common_prefix_len(a: list[int], b: list[int]) -> int: - n = min(len(a), len(b)) - i = 0 - while i < n and a[i] == b[i]: - i += 1 - return i - - -def _output_log_probs(turn: TurnRecord) -> list[float]: - if len(turn.output_log_probs) == len(turn.output_ids): - return list(turn.output_log_probs) - logger.warning( - "[trajectory] turn logprob length mismatch; zeroing output logprobs (%d ids, %d logprobs)", - len(turn.output_ids), - len(turn.output_log_probs), - ) - return [0.0] * len(turn.output_ids) - - -def merge_turns(turns: list[TurnRecord], *, metadata: dict[str, Any] | None = None) -> TokenSegment | None: - """Replay turn records into one linear training segment. - - The first turn's prompt becomes the segment prompt. Later turn prompts are - stitched against ``prompt + response_so_far``. Any new prompt suffix is - non-model context and receives loss mask 0. If a later prompt diverges - inside a previous model output, the retained prefix of that whole output - turn is also masked out, because partial token matches are not a faithful - training target for that turn. + The two kinds are distinguished by whether ``turn`` is set, which reflects + WHERE the message came from: + + * **generated** (``turn is not None``): an assistant message the model + actually generated this turn, fed in via ``record_turn``. ``turn`` holds + its :class:`TurnRecord` -- the prompt/output ids, logprobs and finish + reason that ``get_trajectory`` linearizes into training tokens. + * **routing-only** (``turn is None``): the message came from the prompt, not + from generation, so it only exists to route. This is every + system/user/tool node, AND any assistant we did NOT generate: a foreign + assistant the client replayed in a later prompt, or a prior generated turn + demoted by the rewrite-merge in ``_try_merge_assistant_rewrite``. """ - if not turns: - return None - - prompt_ids = list(turns[0].prompt_ids) - response_ids: list[int] = [] - loss_mask: list[int] = [] - rollout_log_probs: list[float] = [] - output_spans: list[tuple[int, int]] = [] - - for i, turn in enumerate(turns): - if i > 0: - if turn.prompt_ids[: len(prompt_ids)] != prompt_ids: - logger.warning("[trajectory] merge prompt base changed; starting segment from drifted prompt") - prompt_ids = list(turn.prompt_ids) - response_ids = [] - loss_mask = [] - rollout_log_probs = [] - output_spans = [] - else: - prompt_suffix = turn.prompt_ids[len(prompt_ids) :] - matched_len = _common_prefix_len(response_ids, prompt_suffix) - if matched_len < len(response_ids): - logger.warning( - "[trajectory] merge prefix drift; truncating %d unstitched response tokens", - len(response_ids) - matched_len, - ) - for start, end in output_spans: - if start < matched_len < end: - loss_mask[start:matched_len] = [0] * (matched_len - start) - rollout_log_probs[start:matched_len] = [0.0] * (matched_len - start) - response_ids = response_ids[:matched_len] - loss_mask = loss_mask[:matched_len] - rollout_log_probs = rollout_log_probs[:matched_len] - output_spans = [ - (start, min(end, matched_len)) for start, end in output_spans if start < matched_len - ] - - context_tail = prompt_suffix[matched_len:] - response_ids.extend(context_tail) - loss_mask.extend([0] * len(context_tail)) - rollout_log_probs.extend([0.0] * len(context_tail)) - - output_start = len(response_ids) - response_ids.extend(turn.output_ids) - loss_mask.extend([1] * len(turn.output_ids)) - rollout_log_probs.extend(_output_log_probs(turn)) - output_spans.append((output_start, len(response_ids))) - - rollout_log_probs = [logprob if mask else 0.0 for logprob, mask in zip(rollout_log_probs, loss_mask, strict=True)] - - return TokenSegment( - prompt_ids=prompt_ids, - response_ids=response_ids, - loss_mask=loss_mask, - rollout_log_probs=rollout_log_probs, - metadata=dict(metadata or {}), - ) - - -def merge_turn_segments(segments: list[TurnSegment]) -> list[TokenSegment]: - """Merge frozen turn segments and keep every non-empty output.""" - out: list[TokenSegment] = [] - for turn_segment in segments: - token_segment = merge_turns(turn_segment.turns, metadata=turn_segment.metadata) - if token_segment is None: - continue - if token_segment.response_ids: - out.append(token_segment) - return out - - -def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: - """Populate token, mask, response, reward, and status fields from a segment.""" - sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) - sample.response_length = len(segment.response_ids) - sample.loss_mask = list(segment.loss_mask) - sample.rollout_log_probs = list(segment.rollout_log_probs) - sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) - sample.reward = float(reward) - sample.status = Sample.Status.COMPLETED - - -def fan_out_sample_segments( - sample: Sample, - segments: list[TokenSegment], - reward: float, - tokenizer, - *, - metadata: dict[str, Any] | None = None, - rollout_id: int | None = None, -) -> list[Sample]: - """Emit one Sample per segment, splitting reward uniformly across them. - - Sibling samples share ``rollout_id`` so reducers that average by rollout do - not over-count trajectories split by compaction or sub-agent dispatch. + + def __init__( + self, + *, + role: str | None = None, + message: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + parent: MessageNode | None = None, + ) -> None: + self.role = role + self.message = message + self.metadata = dict(metadata or {}) + self.parent: MessageNode | None = parent + self.children: list[MessageNode] = [] + self.turn: TurnRecord | None = None # the generated TurnRecord, else None (routing-only) + self.turn_index: int | None = None + # Shared by sibling leaf paths; the first to reach it trains on it, the rest + # re-emit it as loss_mask=0 context -- so each response is trained exactly once. + self.response_trained: bool = False + + @property + def is_root(self) -> bool: + return self.parent is None + + def add_child(self, child: MessageNode) -> MessageNode: + child.parent = self + self.children.append(child) + return child + + def path_from_root(self) -> list[MessageNode]: + """Ordered list of nodes from the first non-root ancestor down to self.""" + chain: list[MessageNode] = [] + node: MessageNode | None = self + while node is not None and not node.is_root: + chain.append(node) + node = node.parent + chain.reverse() + return chain + + def leaves(self) -> Iterator[MessageNode]: + if not self.children: + yield self + return + for c in self.children: + yield from c.leaves() + + +# =========================================================================== +# drift classification — how an incoming turn's prompt relates to held tokens +# =========================================================================== + + +def _common_prefix_len(a: list[int], b: list[int], chunk: int = 4096) -> int: + limit = min(len(a), len(b)) + matched = 0 + while matched < limit: + chunk_end = min(matched + chunk, limit) + if a[matched:chunk_end] == b[matched:chunk_end]: + matched = chunk_end + else: + while matched < chunk_end and a[matched] == b[matched]: + matched += 1 + return matched + return matched + + +class DriftKind(enum.Enum): + CLEAN = "clean" # drift == 0: prompt_ids exactly extends held tokens; append the tail beyond them + REALIGN = "realign" # drift inside the most-recent response span and short incoming response; replace that span (loss_mask=0) + FORK = "fork" # everything else: close this builder, open a fresh one as a fork + + +# =========================================================================== +# SampleBuilder — accumulates turns into one trainable Sample (fork closes it) +# =========================================================================== + + +class _SampleBuilder: + """Accumulates a chain's turns into the token sequence of one ``Sample``. + + A chain of turns is appended one at a time via :meth:`append_turn`. Ideally + each turn's prompt exactly extends the tokens we already hold, but a replayed + turn rarely re-tokenizes byte-for-byte: TITO round-trips and chat-template + re-rendering both perturb the ids of content we've already seen. The builder + handles this drift in a source-agnostic way, classified by where and how far + the prompt diverges from the held tokens (see :meth:`classify_token_drift`): + + * **CLEAN** -- no drift; append the prompt tail beyond what we hold. + * **REALIGN** -- a short divergence inside the most-recent response span; + overwrite that span from the prompt as loss_mask=0 and keep accumulating. + * **FORK** -- divergence too large or too early to absorb; this builder is + rejected and the caller closes it and opens a fresh one. That boundary is + the "fork". + + Each surviving builder yields one Sample. """ - k = len(segments) - per_segment_reward = float(reward) / max(1, k) - shared_rollout_id = getattr(sample, "index", None) if rollout_id is None else rollout_id - base_metadata = {**(sample.metadata or {}), **(metadata or {})} - - out: list[Sample] = [] - for i, segment in enumerate(segments): - sub = sample if i == 0 else copy.copy(sample) - write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) - sub.rollout_id = shared_rollout_id - sub.metadata = { - **base_metadata, - **(segment.metadata or {}), - "segment_idx": i, - "num_segments": k, + + def __init__(self, fork_threshold: int) -> None: + self._fork_threshold = fork_threshold + self.tokens: list[int] = [] + self.loss_mask: list[int] = [] + self.logprobs: list[float] = [] + self.last_response_start_idx: int | None = None + self.leading_prompt_len: int = 0 + + def classify_token_drift(self, turn: TurnRecord) -> DriftKind: + """Decide how this builder should absorb ``turn``'s prompt. + + The incoming turn's prompt is expected to match the tokens this builder + already holds as an exact prefix. When token drift has occurred -- the + prompt diverges from the held tokens -- we decide whether to REALIGN + (heal a short divergence inside the most-recent response span) or to FORK + (``len(turn.output_ids) >= fork_threshold``, or the divergence sits too + early to absorb). With no drift the turn is handled the CLEAN way -- a + plain prefix extension. + """ + realign_at = _common_prefix_len(self.tokens, turn.prompt_ids) + drift = len(self.tokens) - realign_at + + if drift == 0: + return DriftKind.CLEAN + + # REALIGN only heals drift that falls inside the most-recent response span + # (and is short); divergence anywhere earlier, or an empty builder, forks. + start = self.last_response_start_idx + if start is not None and realign_at >= start and len(turn.output_ids) < self._fork_threshold: + return DriftKind.REALIGN + return DriftKind.FORK + + def append_turn(self, turn: TurnRecord, kind: DriftKind, *, trained: bool = True) -> None: + """Append one turn into this SampleBuilder, branching on ``kind``: for REALIGN + we overwrite the already-saved response span, for CLEAN we just append this + turn's prompt tail.""" + assert kind is not DriftKind.FORK, "append_turn called on a builder that would fork" + + is_first_turn = self.last_response_start_idx is None + + # --- append this turn's prompt tail (loss_mask=0) --- + if kind is DriftKind.REALIGN: + self._align_to_prompt(turn.prompt_ids) # drop the drifted tail, re-append from prompt + else: # CLEAN: held tokens are an exact prefix of prompt_ids; append the tail beyond them + self._append_tokens(turn.prompt_ids[len(self.tokens) :], loss_mask=0) + + # --- append this turn's generated response (loss_mask=1 unless re-emitted as context) --- + self.last_response_start_idx = len(self.tokens) + self._append_tokens( + turn.output_ids, loss_mask=int(trained), logprobs=turn.output_log_probs if trained else None + ) + + if is_first_turn: + self.leading_prompt_len = len(turn.prompt_ids) + + def _align_to_prompt(self, prompt_ids: list[int]) -> None: + """Heal REALIGN drift by overwriting the most-recent response span with + ``prompt_ids`` as loss_mask=0: the drifted tokens carry no signal, and re-appending + from the prompt keeps the builder contiguous. Earlier turns are untouched.""" + response_start = self.last_response_start_idx + tail = prompt_ids[response_start:] + self.tokens[response_start:] = tail + self.loss_mask[response_start:] = [0] * len(tail) + self.logprobs[response_start:] = [0.0] * len(tail) + + def _append_tokens(self, ids: list[int], *, loss_mask: int, logprobs: list[float] | None = None) -> None: + self.tokens.extend(ids) + self.loss_mask.extend([loss_mask] * len(ids)) + self.logprobs.extend(logprobs if logprobs else [0.0] * len(ids)) + + def has_trained_response(self) -> bool: + return any(self.loss_mask[self.leading_prompt_len :]) + + def to_sample( + self, base_sample: Sample, extra_metadata: dict[str, Any] | None, max_sample_tokens: int = 0 + ) -> Sample: + """Emit the accumulated tokens as one ``Sample``, stripping the first-turn + prompt so loss_mask / logprobs cover only the response region.""" + start = self.leading_prompt_len # first-turn prompt stripped; response region starts here + tokens = list(self.tokens) + loss_mask = self.loss_mask + logprobs = self.logprobs + if max_sample_tokens and len(tokens) > max_sample_tokens: + tokens = tokens[:max_sample_tokens] + loss_mask = loss_mask[:max_sample_tokens] + logprobs = logprobs[:max_sample_tokens] + md = dict(extra_metadata or {}) + return Sample( + index=base_sample.index, + group_index=base_sample.group_index, + rollout_id=base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index, + prompt=base_sample.prompt, + label=base_sample.label, + tokens=tokens, + response_length=len(loss_mask) - start, + loss_mask=loss_mask[start:], + rollout_log_probs=logprobs[start:], + reward=0.0, + status=Sample.Status.COMPLETED, + metadata=md, + ) + + +# =========================================================================== +# TrajectoryManager +# =========================================================================== + + +class TrajectoryManager: + def __init__(self, *, fork_threshold_tokens: int | None = None) -> None: + self._fork_threshold: int = 1024 if fork_threshold_tokens is None else fork_threshold_tokens + self._trees: dict[str, MessageNode] = {} + self._turn_count: dict[str, int] = {} + + # -------------------- public ------------------------------------------ + + def has_session(self, sid: str) -> bool: + return sid in self._trees + + def turn_count(self, sid: str) -> int: + return self._turn_count.get(sid, 0) + + def record_turn( + self, + sid: str, + *, + turn: TurnRecord, + prompt_messages: list[dict[str, Any]], + response_message: dict[str, Any] | None, + metadata: dict[str, Any] | None = None, + ) -> None: + if not prompt_messages: + logger.warning("record_turn(sid=%s): empty prompt_messages; skipping", sid) + return + assert not turn.output_log_probs or len(turn.output_log_probs) == len(turn.output_ids), ( + f"turn.output_log_probs length {len(turn.output_log_probs)} != " + f"turn.output_ids length {len(turn.output_ids)}" + ) + + root = self._trees.setdefault(sid, MessageNode()) + + node, depth = self._find_mount_point(root, prompt_messages) + node, depth = self._try_merge_assistant_rewrite(sid, node, prompt_messages, depth) + node = self._mount_prompt_messages(node, prompt_messages[depth:]) + self._attach_assistant_leaf(sid, node, turn=turn, response_message=response_message, metadata=metadata) + + def get_trajectory( + self, + sid: str, + *, + base_sample: Sample, + reward: float = 0.0, + extra_metadata: dict[str, Any] | None = None, + max_sample_tokens: int = 0, + ) -> list[Sample]: + """Linearize this sid's routing tree into vime ``Sample`` objects and + consume the session. + + Each routing leaf yields one or more Samples; ``reward`` is assigned in + full to every emitted Sample (not split across them), so each trained + turn carries the trajectory's outcome reward. The sid is dropped + afterwards, so a second call for the same sid returns ``[]``. + """ + root = self._trees.get(sid) + if root is None: + return [] + + samples: list[Sample] = [] + for routing_leaf in root.leaves(): + if routing_leaf.is_root: + continue + chain = routing_leaf.path_from_root() + samples.extend( + self._chain_to_samples( + chain, base_sample=base_sample, extra_metadata=extra_metadata, max_sample_tokens=max_sample_tokens + ) + ) + + for s in samples: + s.reward = reward + + self._trees.pop(sid, None) + self._turn_count.pop(sid, None) + return samples + + def drop_session(self, sid: str) -> None: + self._trees.pop(sid, None) + self._turn_count.pop(sid, None) + + # -------------------- internals ---------------------------------------- + + def _find_mount_point(self, root: MessageNode, messages: list[dict[str, Any]]) -> tuple[MessageNode, int]: + """Walk down the tree matching each message by role and dict equality (==), + returning the deepest node that still matches and where to mount the rest.""" + node = root + depth = 0 + while depth < len(messages): + msg = messages[depth] + next_child = None + for child in node.children: + if child.role == msg.get("role") and child.message == msg: + next_child = child + break + if next_child is None: + break + node = next_child + depth += 1 + return node, depth + + def _try_merge_assistant_rewrite( + self, + sid: str, + node: MessageNode, + prompt_messages: list[dict[str, Any]], + depth: int, + ) -> tuple[MessageNode, int]: + """Merge a short assistant-rewrite onto its node instead of forking. + + A harness may replay a prior assistant message slightly re-rendered (e.g. + whitespace) in a later prompt. It no longer matches the node we generated, + so it would fork -- stranding the original generated turn as a dead-end + leaf that still emits its own training Sample. Instead we overwrite that + node's message in place and stop training its generated content (demote to + routing-only), so only the live branch trains. This only applies below + ``fork_threshold``: a long abandoned response carries enough real signal + to fork and train standalone. + + Forking is always safe (a rewrite mounts as routing-only); this is purely + a cleanup. So we merge only when the mount point has exactly one assistant + child that is a leaf, generated (``turn`` set), and short (response < + ``fork_threshold``), and fork otherwise, since absorbing destroys a + generated TurnRecord irreversibly. + """ + if self._fork_threshold <= 0: + return node, depth # feature off + if depth >= len(prompt_messages) or prompt_messages[depth].get("role") != "assistant": + return node, depth # genuine non-assistant history fork -> leave it + + asst_children = [c for c in node.children if c.role == "assistant"] + if len(asst_children) != 1: + if len(asst_children) > 1: + logger.warning( + "record_turn(sid=%s turn=%s): %d assistant children at mount " + "point; can't tell which the rewrite targets, so forking.", + sid, + self._turn_count.get(sid, 0) + 1, + len(asst_children), + ) + return node, depth + + rewritten_node = asst_children[0] + if ( + rewritten_node.children + or rewritten_node.turn is None + or len(rewritten_node.turn.output_ids) >= self._fork_threshold + ): + return node, depth + + rewritten_node.metadata["merged_rewrite"] = { + "abandoned_turn_index": rewritten_node.turn_index, + "abandoned_response_tokens": len(rewritten_node.turn.output_ids), } - out.append(sub) - return out + rewritten_node.turn = None + rewritten_node.turn_index = None + rewritten_node.message = prompt_messages[depth] + return rewritten_node, depth + 1 + + def _mount_prompt_messages( + self, + node: MessageNode, + remaining_messages: list[dict[str, Any]], + ) -> MessageNode: + for m in remaining_messages: + node = node.add_child(MessageNode(role=m.get("role"), message=m)) + return node + + def _attach_assistant_leaf( + self, + sid: str, + node: MessageNode, + *, + turn: TurnRecord, + response_message: dict[str, Any] | None, + metadata: dict[str, Any] | None, + ) -> None: + asst = MessageNode( + role="assistant", + message=response_message, + metadata=dict(metadata or {}), + ) + asst.turn = turn + asst.turn_index = self._turn_count.get(sid, 0) + 1 + node.add_child(asst) + self._turn_count[sid] = asst.turn_index + + def _split_chain_into_builders(self, chain: list[MessageNode]) -> list[_SampleBuilder]: + """Pack the chain's generated turns into per-Sample token builders. + + Turns flow into the current builder until one can't extend it as an + exact prefix (re-tokenization drift past what we can drop); that turn + opens a new builder -- a fork. A generated turn shared by sibling leaves + is trained only on the first leaf to claim it; later leaves re-emit it + as loss_mask=0 context so the shared prefix isn't double-counted. + """ + asst_nodes = [n for n in chain if n.role == "assistant" and n.turn is not None] + + builders: list[_SampleBuilder] = [] + for asst_node in asst_nodes: + trained = not asst_node.response_trained + asst_node.response_trained = True + + if not builders or (kind := builders[-1].classify_token_drift(asst_node.turn)) is DriftKind.FORK: + builders.append(_SampleBuilder(self._fork_threshold)) + builders[-1].append_turn(asst_node.turn, DriftKind.CLEAN, trained=trained) + else: + builders[-1].append_turn(asst_node.turn, kind, trained=trained) + return builders + + def _chain_to_samples( + self, + chain: list[MessageNode], + *, + base_sample: Sample, + extra_metadata: dict[str, Any] | None, + max_sample_tokens: int = 0, + ) -> list[Sample]: + + asst_nodes = [n for n in chain if n.role == "assistant" and n.turn is not None] + truncated = bool(asst_nodes) and asst_nodes[-1].turn.finish_reason == "length" + use_tool = any(bool((n.message or {}).get("tool_calls")) for n in asst_nodes) + ill_formed = any(n.turn.ill_formed for n in asst_nodes) + md = { + **(extra_metadata or {}), + "truncated": truncated, + "use_tool": use_tool, + "ill_formed": ill_formed, + } + return [ + builder.to_sample(base_sample, md, max_sample_tokens) + for builder in self._split_chain_into_builders(chain) + if builder.has_trained_response() + ] + + +__all__ = [ + "TrajectoryManager", + "TurnRecord", +] diff --git a/vime/backends/megatron_utils/__init__.py b/vime/backends/megatron_utils/__init__.py index 1497b5a48..feb890366 100644 --- a/vime/backends/megatron_utils/__init__.py +++ b/vime/backends/megatron_utils/__init__.py @@ -2,18 +2,14 @@ import torch -try: - import torch_npu # noqa: F401 -except ImportError: - pass +from vime.platforms import current_platform + +# Load NPU prerequisites before the shared Megatron patches. +current_platform().megatron.bootstrap() -from vime.utils.common import is_npu +from vime.utils import accelerator -if is_npu(): - # MegatronAdaptor must run before Megatron imports so its dummy NPU modules - # are bound in Megatron tensor-parallel modules. - import megatron_adaptor # noqa: F401 - from . import npu_attention_patch # noqa: F401 +accelerator.initialize_accelerator() try: import deep_ep @@ -22,36 +18,35 @@ old_init = deep_ep.Buffer.__init__ def new_init(self, *args, **kwargs): - if torch_memory_saver._impl is not None: - torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(False) - old_init(self, *args, **kwargs) - torch.cuda.synchronize() - if torch_memory_saver._impl is not None: - torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(True) + tms_impl = torch_memory_saver._impl + if tms_impl is None: + return old_init(self, *args, **kwargs) + + cdll = tms_impl._binary_wrapper.cdll + original_interesting_region = cdll.tms_get_interesting_region() + cdll.tms_set_interesting_region(False) + try: + old_init(self, *args, **kwargs) + # DeepEP owns persistent buffers and may initialize them on its + # internal streams. Make their lifetime independent of the TMS + # disabled region before restoring allocation tracking. + # CPU-only imports intentionally have no selected device; explicit + # accelerator requests still fail fast in initialize_accelerator(). + selected_accelerator = accelerator.initialize_accelerator() + if selected_accelerator is not None: + selected_accelerator.synchronize() + else: + # Keep the historical CUDA hook observable for CPU test + # doubles, while ignoring the expected no-CUDA runtime error. + try: + torch.cuda.synchronize() + except RuntimeError: + pass + finally: + cdll.tms_set_interesting_region(original_interesting_region) deep_ep.Buffer.__init__ = new_init except ImportError: logging.warning("deep_ep is not installed, some functionalities may be limited.") -try: - from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import ( - Qwen3VLMoETextRotaryEmbedding, - Qwen3VLTextRotaryEmbedding, - ) - - def patch_rotary_embedding(cls): - _original_forward = cls.forward - - def _patched_forward(self, *args, packed_seq_params=None, **kwargs): - return _original_forward(self, *args, **kwargs) - - cls.forward = _patched_forward - - patch_rotary_embedding(Qwen3VLTextRotaryEmbedding) - patch_rotary_embedding(Qwen3VLMoETextRotaryEmbedding) -except ImportError: - pass - logging.getLogger("megatron").setLevel(logging.WARNING) - -from . import megatron_patch # noqa: F401, E402 diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 68aac0f71..e14682b22 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -1,71 +1,53 @@ import logging import os -import random from argparse import Namespace from contextlib import nullcontext +from datetime import timedelta +from pathlib import Path -import numpy as np import ray import torch import torch.distributed as dist - -from vime.utils.common import is_npu - -if is_npu(): - import importlib - - import megatron_adaptor # noqa: F401 - - importlib.import_module("vime.backends.megatron_utils.npu_attention_patch") - from megatron_adaptor.features_manager.features_manager import FeaturesManager - from megatron_adaptor.utils.args_utils import get_full_args - - def _repatch_megatron_adaptor(args): - """Reapply MegatronAdaptor features after VIME has finalized args.""" - full_args = get_full_args() - for key, value in vars(args).items(): - setattr(full_args, key, value) - FeaturesManager.remove_patches() - FeaturesManager.apply_features_pre_patches(full_args) - FeaturesManager.apply_features_patches(full_args) - - _orig_npu_empty_cache = torch.npu.empty_cache - - def _safe_empty_cache(): - try: - _orig_npu_empty_cache() - except RuntimeError: - pass - - torch.npu.empty_cache = _safe_empty_cache - torch.cuda.empty_cache = _safe_empty_cache from megatron.core import mpu from torch_memory_saver import torch_memory_saver from transformers import AutoConfig, AutoTokenizer +from vime.observability import train_data_utils, train_metric_utils +from vime.observability.logging_utils import init_tracking +from vime.observability.profile_utils import TrainProfiler +from vime.observability.timer import Timer, inverse_timer, timer, with_defer +from vime.platforms import current_platform from vime.ray.train_actor import TrainRayActor -from vime.utils import train_dump_utils +from vime.utils import accelerator from vime.utils.data import process_rollout_data from vime.utils.distributed_utils import get_gloo_group -from vime.utils.logging_utils import init_tracking from vime.utils.memory_utils import clear_memory, print_memory from vime.utils.misc import Box -from vime.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups +from vime.utils.reloadable_process_group import ( + destroy_process_groups, + monkey_patch_torch_dist, + register_default_process_group, + reload_process_groups, +) from vime.utils.routing_replay import RoutingReplay -from vime.utils.timer import Timer, inverse_timer, timer, with_defer from vime.utils.types import RolloutBatch -from ...utils.profile_utils import TrainProfiler from ...utils.tensor_backper import TensorBackuper from .checkpoint import load_checkpoint -from .cp_utils import slice_log_prob_with_cp, slice_with_cp -from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data +from .cp_utils import prepare_routed_experts_for_routing_replay, slice_log_prob_with_cp +from .data import DataIterator, get_data_iterator +from .hf_checkpoint_saver import save_hf_model_to_path from .initialize import init, is_megatron_main_rank -from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values +from .loss import ( + compute_advantages_and_returns, + drain_captured_log_probs, + enable_log_prob_capture, + get_log_probs_and_entropy, + get_values, +) from .model import forward_only, initialize_model_and_optimizer, save, train +from .update_weight import create_weight_updater from .update_weight.common import named_params_and_buffers -from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed -from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor logging.getLogger("megatron").setLevel(logging.WARNING) @@ -87,11 +69,16 @@ def init( monkey_patch_torch_dist() super().init(args, role, with_ref, with_opd_teacher) + # Destroying and recreating WORLD invalidates raw dist.group.WORLD references cached by external code. + # Set VIME_DESTROY_WORLD_PROCESS_GROUP=0 when such references may outlive a train sleep/wake cycle. + if os.getenv("VIME_DESTROY_WORLD_PROCESS_GROUP", "1").lower() not in {"0", "false", "no"}: + register_default_process_group(timeout=timedelta(minutes=args.distributed_timeout_minutes)) + else: + logger.info("Default WORLD process-group destruction is disabled") init(args) - if is_npu(): - _repatch_megatron_adaptor(args) + current_platform().megatron.repatch(args) if is_megatron_main_rank(): init_tracking(args, primary=False, role=role) @@ -106,22 +93,12 @@ def init( dist.barrier(group=get_gloo_group()) - if args.offload_train: - if (x := args.train_memory_margin_bytes) > 0: - logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}") - torch_memory_saver.memory_margin_bytes = x - - tms_region_ctx = None - if args.offload_train and is_npu(): - tms_region_ctx = torch_memory_saver.region(tag="training", enable_cpu_backup=True) - tms_region_ctx.__enter__() - - self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer( - args, role - ) - - if tms_region_ctx is not None: - tms_region_ctx.__exit__(None, None, None) + with current_platform().megatron.training_context(args.offload_train): + self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer( + args, role + ) + if args.offload_train: + current_platform().megatron.initialize_optimizer_state(self.optimizer) vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 if vpp_size > 1: @@ -144,13 +121,8 @@ def init( self.sleep() return start_rollout_id - self.weights_backuper = TensorBackuper.create( - source_getter=lambda: named_params_and_buffers( - self.args, - self.model, - convert_to_global_name=args.megatron_to_hf_mode == "raw", - ), - single_tag=None, + self.weights_backuper = TensorBackuper( + source_getter=lambda: named_params_and_buffers(self.args, self.model), ) self._active_model_tag: str | None = "actor" self.weights_backuper.backup("actor") @@ -171,15 +143,11 @@ def init( if self.args.vocab_size is None: # Prefer HF config vocab_size (which may include model-native padding) - # over tokenizer vocab_size (which may be smaller, e.g. GPT-OSS). + # over tokenizer vocab_size, which may be smaller. hf_vocab = getattr(self.hf_config, "vocab_size", None) self.args.vocab_size = hf_vocab if hf_vocab is not None else self.tokenizer.vocab_size - if self.args.colocate: - update_weight_cls = UpdateWeightFromTensor - else: - update_weight_cls = UpdateWeightFromDistributed - self.weight_updater = update_weight_cls( + self.weight_updater = create_weight_updater( self.args, self.model, weights_getter=lambda: self.weights_backuper.get("actor"), @@ -235,6 +203,15 @@ def wake_up(self) -> None: clear_memory() reload_process_groups() + + if mpu.get_pipeline_model_parallel_world_size() > 2: + # Megatron's patched batched pipeline P2P uses the default WORLD + # group. After reload, PP=4 starts with only the first two stages + # entering batch_isend_irecv(), but PyTorch requires every rank when + # that is the first NCCL operation on a group. Prime WORLD here, + # after the memory saver is resumed, so later stages cannot miss its + # lazy initialization. Sleep still destroys it completely. + dist.barrier(device_ids=[accelerator.current_device()]) if self.role == "actor": self._switch_model("actor") print_memory("after wake_up model") @@ -243,35 +220,32 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: # Fetch data through ray on CPU, not sure if this will be performance bottleneck. # Both first pp stage and the last pp stage will receive the data. rollout_data = process_rollout_data( - self.args, rollout_data_ref, mpu.get_data_parallel_rank(with_context_parallel=False), mpu.get_data_parallel_world_size(with_context_parallel=False), ) # TODO: this is ugly, move to somewhere else? # move tokens to GPU in advance - device = torch.npu.current_device() if is_npu() else torch.cuda.current_device() - rollout_data["tokens"] = [torch.tensor(t, dtype=torch.long, device=device) for t in rollout_data["tokens"]] + device = accelerator.current_device() + rollout_data["tokens"] = [ + t.to(device=device, dtype=torch.long, non_blocking=True) for t in rollout_data["tokens"] + ] rollout_data["loss_masks"] = [ - torch.tensor(t, dtype=torch.int, device=device) for t in rollout_data["loss_masks"] + t.to(device=device, dtype=torch.int, non_blocking=True) for t in rollout_data["loss_masks"] ] if "rollout_mask_sums" in rollout_data: # Promote precomputed per-rollout mask totals to GPU tensors here # (matching loss_masks) so the loss reducer can just divide. - rollout_data["rollout_mask_sums"] = torch.tensor( - rollout_data["rollout_mask_sums"], dtype=torch.float32, device=torch.cuda.current_device() + rollout_data["rollout_mask_sums"] = rollout_data["rollout_mask_sums"].to( + device=device, dtype=torch.float32, non_blocking=True ) if "multimodal_train_inputs" in rollout_data: # Move multimodal training tensors to GPU in advance rollout_data["multimodal_train_inputs"] = [ ( { - key: ( - torch.from_numpy(v.copy()).to(device=device) - if isinstance(v, np.ndarray) - else v.to(device=device) - ) - for key, v in mm_dict.items() + key: value.to(device=device, non_blocking=True) if isinstance(value, torch.Tensor) else value + for key, value in mm_dict.items() } if mm_dict is not None else None @@ -279,44 +253,22 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: for mm_dict in rollout_data["multimodal_train_inputs"] ] - if self.args.qkv_format == "bshd": - # TODO: micro-batch wise dynamic, possibly move to @data.py:get_data_iterator - max_seq_len = max(rollout_data["total_lengths"]) - - # pad to reduce memory fragmentation and maybe make the computation faster - pad_size = mpu.get_tensor_model_parallel_world_size() * self.args.data_pad_size_multiplier - max_seq_len = (max_seq_len + pad_size - 1) // pad_size * pad_size - - rollout_data["max_seq_lens"] = [max_seq_len] * len(rollout_data["tokens"]) - for key in ["rollout_log_probs", "teacher_log_probs"]: if key not in rollout_data: continue rollout_data[key] = [ - torch.tensor( - slice_log_prob_with_cp( - log_prob, - total_length, - response_length, - self.args.qkv_format, - rollout_data["max_seq_lens"][i] if self.args.qkv_format == "bshd" else None, - ), - device=torch.cuda.current_device(), + slice_log_prob_with_cp(log_prob, total_length, response_length).to( + device=device, dtype=torch.float32, + non_blocking=True, ) - for i, (log_prob, total_length, response_length) in enumerate( - zip( - rollout_data[key], - rollout_data["total_lengths"], - rollout_data["response_lengths"], - strict=False, - ) + for log_prob, total_length, response_length in zip( + rollout_data[key], + rollout_data["total_lengths"], + rollout_data["response_lengths"], + strict=False, ) ] - if "rollout_routed_experts" in rollout_data: - rollout_data["rollout_routed_experts"] = [ - torch.from_numpy(r) for r in rollout_data["rollout_routed_experts"] - ] return rollout_data def _switch_model(self, target_tag: str) -> None: @@ -339,45 +291,16 @@ def fill_routing_replay(self, data_iterator, num_microbatches, rollout_data): for iterator in data_iterator: iterator.reset() - tp_rank = mpu.get_tensor_model_parallel_rank() - tp_size = mpu.get_tensor_model_parallel_world_size() - - def pad_func(experts, pad): - _, num_layers, topk = experts.shape - pad = ( - torch.arange( - pad * num_layers * topk, - device=experts.device, - dtype=experts.dtype, - ).reshape((pad, num_layers, topk)) - % self.args.num_experts - ) - return torch.cat([experts, pad], dim=0) - for _ in range(sum(num_microbatches)): batch = data_iterator[0].get_next(["rollout_routed_experts", "tokens"]) - rollout_routed_experts = batch["rollout_routed_experts"] - tokens = batch["tokens"] - assert len(rollout_routed_experts) == len(tokens) - for a, b in zip(rollout_routed_experts, tokens, strict=False): - assert a.shape[0] == b.shape[0] - 1, f"{a.shape}, {b.shape}" - - # We need to pad the experts to the last token. We won't calculate loss on this token so this should be fine. - # TODO: fuse this padding with the following slice_with_cp to reduce memory copy. - rollout_routed_experts = [pad_func(r, 1) for r in rollout_routed_experts] - # TODO: maybe extract a common process function for here and get_batch? - rollout_routed_experts = [slice_with_cp(r, pad_func) for r in rollout_routed_experts] - rollout_routed_experts = torch.cat(rollout_routed_experts, dim=0) - pad_size = mpu.get_tensor_model_parallel_world_size() * self.args.data_pad_size_multiplier - pad = (pad_size - rollout_routed_experts.size(0) % pad_size) % pad_size - if pad != 0: - rollout_routed_experts = pad_func(rollout_routed_experts, pad) - - if self.args.sequence_parallel: - seqlen = rollout_routed_experts.size(0) - assert seqlen % tp_size == 0 - start, end = seqlen // tp_size * tp_rank, seqlen // tp_size * (tp_rank + 1) - rollout_routed_experts = rollout_routed_experts[start:end] + rollout_routed_experts = prepare_routed_experts_for_routing_replay( + batch["rollout_routed_experts"], + batch["tokens"], + num_experts=self.args.num_experts, + data_pad_size_multiplier=self.args.data_pad_size_multiplier, + sequence_parallel=self.args.sequence_parallel, + allgather_cp=self.args.allgather_cp, + ) routing_replay_offset = 0 for vp_stage, model in enumerate(self.model): @@ -409,7 +332,6 @@ def compute_log_prob( num_microbatches: list[int], store_prefix: str = "", ) -> dict[str, list[torch.Tensor]]: - with timer(f"{store_prefix}log_probs"): return forward_only( get_log_probs_and_entropy, @@ -418,6 +340,7 @@ def compute_log_prob( data_iterator, num_microbatches, store_prefix=store_prefix, + use_rollout_top_p_replay=True, ) def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): @@ -516,7 +439,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data and not self.args.use_critic and not self.args.keep_old_actor and not self.args.use_opd - and not self.args.use_routing_replay + and (not self.args.use_routing_replay or self.args.use_rollout_routing_replay) and self.args.advantage_estimator != "gspo" ) if ( @@ -554,7 +477,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data if self.rollout_data_postprocess is not None: self.rollout_data_postprocess(self.args, rollout_id, rollout_data) - log_rollout_data( + train_metric_utils.log_rollout_data( rollout_id, self.args, rollout_data, @@ -563,6 +486,13 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data # Train if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" + # When dumping train debug data but the actor log_probs were not + # recomputed separately (can_reuse_log_probs_in_loss / use_rollout_logprobs), + # snapshot them from the training forward so the dump still carries + # per-sample log_probs — at no extra forward pass. + capture_log_probs = self.args.save_debug_train_data is not None and "log_probs" not in rollout_data + if capture_log_probs: + enable_log_prob_capture() with timer("actor_train"): train( rollout_id, @@ -573,10 +503,18 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data num_microbatches, global_batch_sizes, ) + if capture_log_probs: + captured = drain_captured_log_probs() + # `captured` is non-empty only on the last PP stage running a loss + # that snapshots log_probs (policy_loss), and then covers every + # local sample. Key it by this rank's `partition` to land in local + # sample order; skip otherwise (nothing to place). + if captured: + rollout_data["log_probs"] = [captured[pos] for pos in rollout_data["partition"]] self.prof.step(rollout_id=rollout_id) - train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data) + train_data_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data) if self.args.use_routing_replay: RoutingReplay.clear_all() @@ -595,7 +533,11 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") - log_perf_data(rollout_id, self.args) + train_metric_utils.log_perf_data( + rollout_id, + self.args, + extra_metrics=self.weight_updater.pop_metrics(), + ) @timer def save_model(self, rollout_id: int, force_sync: bool = False) -> None: @@ -617,9 +559,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: maybe_finalize_async_save(blocking=True) if self.args.save_hf is not None and self.role == "actor": - from vime.backends.megatron_utils.model import save_hf_model - - save_hf_model(self.args, rollout_id, self.model) + save_hf_model_to_path(self.args, Path(self.args.save_hf.format(rollout_id=rollout_id)), self.model) if self.args.offload_train: self.sleep() @@ -634,12 +574,22 @@ def update_weights(self) -> None: ray.get(self.rollout_manager.recover_updatable_engines.remote()) dist.barrier(group=get_gloo_group()) - rollout_engines, rollout_engine_lock, num_new_engines, engine_gpu_counts, engine_gpu_offsets = ray.get( - self.rollout_manager.get_updatable_engines_and_lock.remote() - ) + ( + rollout_engines, + rollout_engine_lock, + num_new_engines, + engine_gpu_counts, + engine_gpu_offsets, + engine_parallel_configs, + ) = ray.get(self.rollout_manager.get_updatable_engines_and_lock.remote()) reconnect_rollout_engines = self.args.offload_train and self.args.use_critic and not self.args.colocate + if not rollout_engines and not reconnect_rollout_engines: + if dist.get_rank() == 0: + logger.info("No updatable VLLM engines are running; skip weight update.") + return + if reconnect_rollout_engines: self.wake_up() elif self.args.offload_train: @@ -651,24 +601,26 @@ def update_weights(self) -> None: rollout_engine_lock, engine_gpu_counts=engine_gpu_counts, engine_gpu_offsets=engine_gpu_offsets, + engine_parallel_configs=engine_parallel_configs, ) dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_updatable_num_new_engines.remote()) - with torch_memory_saver.disable() if (self.args.offload_train and not is_npu()) else nullcontext(): + with ( + torch_memory_saver.disable() + if (self.args.offload_train and not current_platform().is_npu) + else nullcontext() + ): + if self.args.dspark_enabled and self.args.offload_train: + backup = self.weights_backuper.get("actor") + for name, param in named_params_and_buffers(self.args, self.model): + if ".draft_model." in name: + param.data = backup[name].to(param.device) print_memory("before update_weights") self.weight_updater.update_weights() print_memory("after update_weights") - if self.args.ci_test and len(rollout_engines) > 0: - engine = random.choice(rollout_engines) - engine_version = ray.get(engine.get_weight_version.remote()) - if str(engine_version) != str(self.weight_updater.weight_version): - raise RuntimeError( - f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}" - ) - if getattr(self.args, "keep_old_actor", False): if self.args.update_weights_interval == 1: logger.info("updating model queue: rollout_actor -> old_actor, actor -> rollout_actor") @@ -686,18 +638,21 @@ def update_weights(self) -> None: destroy_process_groups() def load_other_checkpoint(self, model_tag: str, path: str) -> None: - old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune + old_args = ( + self.args.load, + self.args.no_load_optim, + self.args.no_load_rng, + self.args.finetune, + self.args.ckpt_step, + ) self.args.load = path self.args.no_load_optim = True self.args.no_load_rng = True self.args.finetune = True - old_ckpt_step = None if model_tag == "ref" and self.args.ref_ckpt_step is not None: - old_ckpt_step = self.args.ckpt_step self.args.ckpt_step = self.args.ref_ckpt_step elif model_tag == "teacher" and self.args.opd_teacher_ckpt_step is not None: - old_ckpt_step = self.args.ckpt_step self.args.ckpt_step = self.args.opd_teacher_ckpt_step _, _ = load_checkpoint( @@ -705,12 +660,14 @@ def load_other_checkpoint(self, model_tag: str, path: str) -> None: None, None, checkpointing_context={}, - skip_load_to_model_and_opt=False, ) - self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args - - if old_ckpt_step is not None: - self.args.ckpt_step = old_ckpt_step + ( + self.args.load, + self.args.no_load_optim, + self.args.no_load_rng, + self.args.finetune, + self.args.ckpt_step, + ) = old_args self.weights_backuper.backup(model_tag) self._active_model_tag = model_tag diff --git a/vime/backends/megatron_utils/alignment/__init__.py b/vime/backends/megatron_utils/alignment/__init__.py new file mode 100644 index 000000000..682d8170b --- /dev/null +++ b/vime/backends/megatron_utils/alignment/__init__.py @@ -0,0 +1 @@ +"""Megatron train-side train/rollout numerical-alignment helpers.""" diff --git a/vime/backends/megatron_utils/alignment/deepgemm_forward.py b/vime/backends/megatron_utils/alignment/deepgemm_forward.py new file mode 100644 index 000000000..34f89b68b --- /dev/null +++ b/vime/backends/megatron_utils/alignment/deepgemm_forward.py @@ -0,0 +1,1185 @@ +"""Use VLLM's block-FP8 DeepGEMM result in selected Megatron linears. + +This is an opt-in numerical-alignment hook. It replaces the selected +Transformer Engine linear with a custom autograd function: + +* forward: VLLM-style block-FP8 DeepGEMM; +* backward: explicit BF16 GEMMs for dgrad/wgrad plus analytic norm gradients + for fused LayerNorm/RMSNorm linears. + +The first implementation requires tensor parallel size one. This makes each +target a full matrix, matching VLLM's dense-TP1 execution and avoiding +row-parallel partial-sum rounding as a confounder. +""" + +from __future__ import annotations + +import logging +import os +import re +import types +from collections.abc import Iterable + +import torch +from megatron.core import parallel_state + +logger = logging.getLogger(__name__) + +try: + from megatron.core.extensions.transformer_engine import te_general_gemm +except ImportError: + te_general_gemm = None + + +def router_gating_linear_backward(inp, weight, grad_output, router_dtype): + """Compute router-linear dgrad/wgrad without evaluating its forward. + + Re-homed into vime so the GLM-5 alignment path does not depend on this + helper being present in Megatron ``moe_utils`` (it only exists on newer + Megatron commits). Behaviour matches that upstream helper. + """ + input_dtype = inp.dtype + weight_dtype = weight.dtype + inp_shape = inp.shape + inp_2d = inp.reshape(-1, inp_shape[-1]) + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + if ( + te_general_gemm is not None + and router_dtype != torch.float64 + and inp_2d.is_cuda + and weight.is_cuda + and grad_2d.is_cuda + ): + grad_input = te_general_gemm(weight.to(router_dtype), grad_2d, router_dtype, layout="NN", grad=True)[0].to( + input_dtype + ) + grad_weight = te_general_gemm(inp_2d.to(router_dtype), grad_2d, router_dtype, layout="NT", grad=True)[0].to( + weight_dtype + ) + else: + grad_input = torch.mm(grad_2d, weight.to(router_dtype)).to(input_dtype) + grad_weight = torch.mm(grad_2d.t(), inp_2d.to(router_dtype)).to(weight_dtype) + + return grad_input.reshape(*inp_shape), grad_weight + + +def _vllm_silu_and_mul(input_: torch.Tensor, output: torch.Tensor) -> None: + torch.ops._C.silu_and_mul(output, input_) + + +_LAYER_PATH_RE = re.compile(r"^(?P(?:.*\.)?decoder\.layers\.(?P\d+))\.") +_DEFAULT_TARGET_SUFFIXES = ( + "self_attention.linear_q_down_proj", + "self_attention.linear_q_up_proj", + "self_attention.linear_kv_down_proj", + "self_attention.linear_kv_up_proj", + "self_attention.linear_proj", + "self_attention.wq_b", + "self_attention.wk", + "mlp.linear_fc1", + "mlp.linear_fc2", + "mlp.shared_experts.linear_fc1", + "mlp.shared_experts.linear_fc2", +) +_SUPPORTED_TE_CLASS_NAMES = { + "TELinear", + "TEColumnParallelLinear", + "TERowParallelLinear", + "TELayerNormLinear", + "TELayerNormColumnParallelLinear", +} + + +def _should_log_deepgemm_summary() -> bool: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() == 0 + return True + + +def _format_int_ranges(values: Iterable) -> str: + sorted_values = sorted({int(value) for value in values}) + if not sorted_values: + return "[]" + + ranges = [] + start = prev = sorted_values[0] + for value in sorted_values[1:]: + if value == prev + 1: + prev = value + continue + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + start = prev = value + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + return ",".join(ranges) + + +def _deepgemm_bf16_gemm_nn(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + """Compute ``lhs @ rhs`` using DeepGEMM BF16 when available.""" + if lhs.ndim != 2 or rhs.ndim != 2 or lhs.shape[1] != rhs.shape[0]: + raise RuntimeError(f"BF16 NN GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if lhs.is_cuda and rhs.is_cuda and lhs.dtype == torch.bfloat16 and rhs.dtype == torch.bfloat16: + import deep_gemm + + out = torch.empty((lhs.shape[0], rhs.shape[1]), device=lhs.device, dtype=torch.bfloat16) + deep_gemm.bf16_gemm_nn(lhs.contiguous(), rhs.contiguous(), out) + return out + return lhs.matmul(rhs).to(lhs.dtype) + + +def _deepgemm_bf16_gemm_nt(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + """Compute ``lhs @ rhs.T`` using DeepGEMM BF16 when available.""" + if lhs.ndim != 2 or rhs.ndim != 2 or lhs.shape[1] != rhs.shape[1]: + raise RuntimeError(f"BF16 NT GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if lhs.is_cuda and rhs.is_cuda and lhs.dtype == torch.bfloat16 and rhs.dtype == torch.bfloat16: + import deep_gemm + + out = torch.empty((lhs.shape[0], rhs.shape[0]), device=lhs.device, dtype=torch.bfloat16) + deep_gemm.bf16_gemm_nt(lhs.contiguous(), rhs.contiguous(), out) + return out + return lhs.matmul(rhs.transpose(0, 1)).to(lhs.dtype) + + +def _deepgemm_bf16_gemm_tn(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + """Compute ``lhs.T @ rhs`` using DeepGEMM BF16 when available.""" + if lhs.ndim != 2 or rhs.ndim != 2 or lhs.shape[0] != rhs.shape[0]: + raise RuntimeError(f"BF16 TN GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if lhs.is_cuda and rhs.is_cuda and lhs.dtype == torch.bfloat16 and rhs.dtype == torch.bfloat16: + import deep_gemm + + out = torch.empty((lhs.shape[1], rhs.shape[1]), device=lhs.device, dtype=torch.bfloat16) + deep_gemm.bf16_gemm_tn(lhs.contiguous(), rhs.contiguous(), out) + return out + return lhs.transpose(0, 1).matmul(rhs).to(lhs.dtype) + + +def _sum_to_parameter_dtype(value: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + return value.to(dtype=reference.dtype) + + +def _effective_norm_weight(norm_weight: torch.Tensor, zero_centered_gamma: bool) -> torch.Tensor: + weight = norm_weight.float() + if zero_centered_gamma: + weight = weight + 1.0 + return weight + + +def _norm_forward( + input_: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor | None, + *, + normalization: str, + eps: float, + zero_centered_gamma: bool, +) -> torch.Tensor: + if normalization == "RMSNorm" and os.environ.get("MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS", "0") == "1": + if norm_bias is not None: + raise RuntimeError("VLLM RMSNorm alignment does not support a norm bias") + from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant + + weight = norm_weight + if zero_centered_gamma: + weight = (norm_weight.float() + 1.0).to(dtype=norm_weight.dtype) + return rms_norm_batch_invariant(input_, weight, eps) + + x = input_.float() + weight = _effective_norm_weight(norm_weight, zero_centered_gamma) + if normalization == "RMSNorm": + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + output = x * rstd * weight + elif normalization == "LayerNorm": + mean = x.mean(dim=-1, keepdim=True) + centered = x - mean + rstd = torch.rsqrt(centered.pow(2).mean(dim=-1, keepdim=True) + eps) + output = centered * rstd * weight + if norm_bias is not None: + output = output + norm_bias.float() + else: + raise RuntimeError(f"Unsupported fused norm type for DeepGEMM wrapper: {normalization}") + return output.to(input_.dtype) + + +def _norm_backward( + grad_output: torch.Tensor, + input_: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor | None, + *, + normalization: str, + eps: float, + zero_centered_gamma: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + x = input_.float() + grad = grad_output.float() + weight = _effective_norm_weight(norm_weight, zero_centered_gamma) + reduce_dims = tuple(range(grad.ndim - 1)) + + if normalization == "RMSNorm": + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + xhat = x * rstd + grad_norm_weight = (grad * xhat).sum(dim=reduce_dims) + scaled_grad = grad * weight + mean_scaled_grad_x = (scaled_grad * x).mean(dim=-1, keepdim=True) + grad_input = scaled_grad * rstd - x * mean_scaled_grad_x * rstd.pow(3) + grad_norm_bias = None + elif normalization == "LayerNorm": + mean = x.mean(dim=-1, keepdim=True) + centered = x - mean + rstd = torch.rsqrt(centered.pow(2).mean(dim=-1, keepdim=True) + eps) + xhat = centered * rstd + grad_norm_weight = (grad * xhat).sum(dim=reduce_dims) + scaled_grad = grad * weight + hidden = x.shape[-1] + grad_input = ( + scaled_grad + - scaled_grad.mean(dim=-1, keepdim=True) + - xhat * (scaled_grad * xhat).mean(dim=-1, keepdim=True) + ) * rstd + if hidden <= 0: + raise RuntimeError("LayerNorm hidden size must be positive") + grad_norm_bias = grad.sum(dim=reduce_dims) if norm_bias is not None else None + else: + raise RuntimeError(f"Unsupported fused norm type for DeepGEMM wrapper: {normalization}") + + return ( + grad_input.to(dtype=input_.dtype), + _sum_to_parameter_dtype(grad_norm_weight, norm_weight), + None if grad_norm_bias is None else _sum_to_parameter_dtype(grad_norm_bias, norm_bias), + ) + + +class _VLLMRMSNormWithAnalyticBackward(torch.autograd.Function): + """VLLM RMSNorm forward with a single-pass analytic RMSNorm backward.""" + + @staticmethod + def forward( + ctx, + visible_input: torch.Tensor, + backward_input: torch.Tensor, + weight: torch.Tensor, + eps: float, + use_fp32_sum: bool, + output_dtype: torch.dtype, + ) -> torch.Tensor: + ctx.eps = eps + ctx.save_for_backward(backward_input, weight) + if use_fp32_sum: + return _vllm_rmsnorm_from_fp32_sum( + visible_input, + weight, + eps, + output_dtype, + ) + return _vllm_batch_invariant_rmsnorm(visible_input, weight, eps) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + backward_input, weight = ctx.saved_tensors + grad_input, grad_weight, _ = _norm_backward( + grad_output, + backward_input, + weight, + None, + normalization="RMSNorm", + eps=ctx.eps, + zero_centered_gamma=False, + ) + if not ctx.needs_input_grad[1]: + grad_input = None + if not ctx.needs_input_grad[2]: + grad_weight = None + return None, grad_input, grad_weight, None, None, None + + +class _VLLMRouterGEMMWithMegatronBackward(torch.autograd.Function): + """Native VLLM GateLinear forward with Megatron backward.""" + + @staticmethod + def forward( + ctx, + value: torch.Tensor, + weight: torch.Tensor, + gate_linear: torch.nn.Module, + ) -> torch.Tensor: + original_shape = value.shape + flat_value = value.reshape(-1, original_shape[-1]) + output = gate_linear(flat_value) + if isinstance(output, tuple): + output, output_bias = output + if output_bias is not None: + raise RuntimeError("VLLM router GateLinear unexpectedly returned a bias") + output = output.view(*original_shape[:-1], -1) + ctx.save_for_backward(value, weight) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + value, weight = ctx.saved_tensors + grad_input, grad_weight = router_gating_linear_backward( + value, + weight, + grad_output, + torch.float32, + ) + if not ctx.needs_input_grad[0]: + grad_input = None + if not ctx.needs_input_grad[1]: + grad_weight = None + return grad_input, grad_weight, None + + +class _VLLMSwiGLUWithAnalyticBackward(torch.autograd.Function): + """VLLM fused SwiGLU forward with an analytic FP32 backward.""" + + @staticmethod + def forward(ctx, value: torch.Tensor) -> torch.Tensor: + output_shape = (*value.shape[:-1], value.shape[-1] // 2) + output = torch.empty(output_shape, dtype=value.dtype, device=value.device) + _vllm_silu_and_mul( + value.contiguous().view(-1, value.shape[-1]), + output.view(-1, output.shape[-1]), + ) + ctx.save_for_backward(value) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if not ctx.needs_input_grad[0]: + return None + (value,) = ctx.saved_tensors + gate, up = value.chunk(2, dim=-1) + gate_fp32 = gate.float() + up_fp32 = up.float() + grad_fp32 = grad_output.float() + sigmoid = torch.sigmoid(gate_fp32) + silu_gate = gate_fp32 * sigmoid + grad_gate = grad_fp32 * up_fp32 * sigmoid * (1.0 + gate_fp32 * (1.0 - sigmoid)) + grad_up = grad_fp32 * silu_gate + return torch.cat((grad_gate, grad_up), dim=-1).to(value.dtype) + + +class _DeepGEMMLinearWithBF16Backward(torch.autograd.Function): + """FP8 DeepGEMM forward with BF16 GEMM dgrad/wgrad.""" + + @staticmethod + def forward( + ctx, + input_: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + output = _deepgemm_linear(input_, weight) + ctx.input_shape = tuple(input_.shape) + ctx.save_for_backward(input_, weight) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + input_, weight = ctx.saved_tensors + grad_2d = grad_output.contiguous().view(-1, grad_output.shape[-1]).to(dtype=input_.dtype) + input_2d = input_.contiguous().view(-1, input_.shape[-1]) + + grad_input = _deepgemm_bf16_gemm_nn(grad_2d, weight).view(ctx.input_shape) if ctx.needs_input_grad[0] else None + grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = _sum_to_parameter_dtype( + _deepgemm_bf16_gemm_tn(grad_2d, input_2d), + weight, + ) + return grad_input, grad_weight + + +class _DeepGEMMLayerNormLinearWithBF16Backward(torch.autograd.Function): + """FP8 DeepGEMM fused norm+linear forward with BF16 GEMM backward.""" + + @staticmethod + def forward( + ctx, + input_: torch.Tensor, + weight: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor | None, + normalization: str, + eps: float, + zero_centered_gamma: bool, + ) -> torch.Tensor: + normalized = _norm_forward( + input_, + norm_weight, + norm_bias, + normalization=normalization, + eps=eps, + zero_centered_gamma=zero_centered_gamma, + ) + output = _deepgemm_linear(normalized, weight) + ctx.input_shape = tuple(input_.shape) + ctx.normalization = normalization + ctx.eps = eps + ctx.zero_centered_gamma = zero_centered_gamma + ctx.has_norm_bias = norm_bias is not None + tensors = (input_, normalized, weight, norm_weight) + if norm_bias is not None: + tensors = (*tensors, norm_bias) + ctx.save_for_backward(*tensors) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + saved = ctx.saved_tensors + input_, normalized, weight, norm_weight = saved[:4] + norm_bias = saved[4] if ctx.has_norm_bias else None + + grad_2d = grad_output.contiguous().view(-1, grad_output.shape[-1]).to(dtype=normalized.dtype) + normalized_2d = normalized.contiguous().view(-1, normalized.shape[-1]) + grad_normalized = _deepgemm_bf16_gemm_nn(grad_2d, weight).view(ctx.input_shape) + grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = _sum_to_parameter_dtype( + _deepgemm_bf16_gemm_tn(grad_2d, normalized_2d), + weight, + ) + grad_input, grad_norm_weight, grad_norm_bias = _norm_backward( + grad_normalized, + input_, + norm_weight, + norm_bias, + normalization=ctx.normalization, + eps=ctx.eps, + zero_centered_gamma=ctx.zero_centered_gamma, + ) + return ( + grad_input if ctx.needs_input_grad[0] else None, + grad_weight, + grad_norm_weight if ctx.needs_input_grad[2] else None, + grad_norm_bias if ctx.needs_input_grad[3] else None, + None, + None, + None, + ) + + +def _deepgemm_linear( + input_: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + import deep_gemm + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, + requant_weight_ue8m0_inplace, + ) + from vllm.utils import deep_gemm as vllm_deep_gemm + + from vime.backends.megatron_utils.alignment.deepgemm_moe_forward import _configure_batch_invariant + from vime.backends.megatron_utils.kernels.fp8_kernel import blockwise_cast_to_fp8_triton + + # Megatron actors are separate processes from VLLM workers. Environment + # propagation alone does not call DeepGEMM's process-local setter. + _configure_batch_invariant(deep_gemm) + + if input_.dtype != torch.bfloat16: + raise RuntimeError(f"DeepGEMM alignment forward requires BF16 input, got {input_.dtype}") + if weight.dtype != torch.bfloat16: + raise RuntimeError(f"DeepGEMM alignment forward requires BF16 weight, got {weight.dtype}") + if input_.shape[-1] != weight.shape[-1]: + raise RuntimeError(f"DeepGEMM input/weight K mismatch: {input_.shape[-1]} != {weight.shape[-1]}") + + with torch.no_grad(): + use_ue8m0 = vllm_deep_gemm.is_deep_gemm_e8m0_used() + if use_ue8m0: + # Native VLLM quantizes the BF16 weight to block FP8 and then + # requantizes it onto UE8M0 power-of-two scales before dispatch. + # Apply the same two helpers to the Megatron weight. + qweight, weight_scale = vllm_deep_gemm.per_block_cast_to_fp8( + weight.detach().contiguous(), block_size=(128, 128) + ) + requant_weight_ue8m0_inplace(qweight, weight_scale, (128, 128)) + else: + # Hopper (sm90): FP32 block scales. Unchanged. + qweight, weight_scale = blockwise_cast_to_fp8_triton(weight.detach().contiguous(), (128, 128)) + qinput, input_scale = per_token_group_quant_fp8( + input_.detach().contiguous(), + 128, + column_major_scales=True, + tma_aligned_scales=True, + use_ue8m0=use_ue8m0, + ) + output = torch.empty( + (*input_.shape[:-1], weight.shape[0]), + dtype=torch.bfloat16, + device=input_.device, + ) + vllm_deep_gemm.fp8_gemm_nt( + (qinput.reshape(-1, qinput.shape[-1]), input_scale), + (qweight, weight_scale), + output.reshape(-1, output.shape[-1]), + is_deep_gemm_e8m0_used=use_ue8m0, + ) + return output + + +def _validate_custom_backward(module: torch.nn.Module) -> None: + config = getattr(module, "config", None) + if config is not None and getattr(config, "delay_wgrad_compute", False): + raise RuntimeError( + "DeepGEMM custom backward does not support Megatron delay_wgrad_compute; " "disable delay_wgrad_compute." + ) + + +def _wrap_te_linear(module: torch.nn.Module, module_name: str) -> bool: + if getattr(module, "_vime_deepgemm_forward_wrapped", False): + return False + + class_name = type(module).__name__ + if class_name not in _SUPPORTED_TE_CLASS_NAMES: + raise RuntimeError(f"DeepGEMM target {module_name} has unsupported module class {class_name}") + if getattr(module, "use_bias", False): + raise RuntimeError(f"DeepGEMM alignment probe currently requires bias-free linears: {module_name}") + + _validate_custom_backward(module) + is_layernorm_linear = "LayerNorm" in class_name + + def deepgemm_forward(self, input_): + input_already_normalized = bool(getattr(self, "_deepgemm_input_already_normalized", False)) + if is_layernorm_linear and not input_already_normalized: + config = getattr(self, "config", None) + if config is None: + raise RuntimeError(f"DeepGEMM custom backward for {module_name} requires module.config") + norm_weight = getattr(self, "layer_norm_weight", None) + if not isinstance(norm_weight, torch.Tensor): + raise RuntimeError(f"{module_name} is missing layer_norm_weight") + norm_bias = getattr(self, "layer_norm_bias", None) + if norm_bias is not None and not isinstance(norm_bias, torch.Tensor): + raise RuntimeError(f"{module_name}.layer_norm_bias is not a Tensor") + output = _DeepGEMMLayerNormLinearWithBF16Backward.apply( + input_, + self.weight, + norm_weight, + norm_bias, + getattr(config, "normalization", "RMSNorm"), + float(getattr(config, "layernorm_epsilon", 1e-5)), + bool(getattr(config, "layernorm_zero_centered_gamma", False)), + ) + return output, None + + output = _DeepGEMMLinearWithBF16Backward.apply( + input_, + self.weight, + ) + return output, None + + module.forward = types.MethodType(deepgemm_forward, module) + module._vime_deepgemm_forward_wrapped = True + module._vime_deepgemm_module_name = module_name + return True + + +def _get_global_layer_index(model_chunk: torch.nn.Module, module_name: str) -> int | None: + match = _LAYER_PATH_RE.search(module_name) + if match is None: + return None + layer_path = match.group("layer_path") + layer = model_chunk.get_submodule(layer_path) + layer_number = getattr(layer, "layer_number", None) + if layer_number is None: + raise RuntimeError( + "DeepGEMM requires TransformerLayer.layer_number to select global pipeline " + f"layers, but {layer_path} (from {module_name}) has no layer_number" + ) + return int(layer_number) - 1 + + +def _as_set(values: Iterable | None, default: Iterable) -> set: + return set(default if values is None else values) + + +def enable_deepgemm_forward(args, model, store_prefix: str) -> None: + """Install the forward-value replacement on selected full-matrix TE linears.""" + del store_prefix + if parallel_state.get_tensor_model_parallel_world_size() != 1: + raise RuntimeError("The initial DeepGEMM alignment probe requires tensor model parallel size 1") + + target_layers = _as_set(args.megatron_deepgemm_forward_layers, ()) + if not target_layers: + raise RuntimeError("--megatron-deepgemm-forward-layers must select at least one layer") + target_suffixes = _as_set( + args.megatron_deepgemm_forward_modules, + _DEFAULT_TARGET_SUFFIXES, + ) + if isinstance(model, torch.nn.Module): + model_chunks = [model] + else: + try: + model_chunks = list(model) + except TypeError: + # Hook-composition tests may use an opaque sentinel while + # monkeypatching the concrete installers. + return + + wrapped = [] + for model_chunk in model_chunks: + for name, module in model_chunk.named_modules(): + if not any(name.endswith(suffix) for suffix in target_suffixes): + continue + global_layer_index = _get_global_layer_index(model_chunk, name) + if global_layer_index not in target_layers: + continue + if _wrap_te_linear(module, name): + wrapped.append(name) + + if wrapped and _should_log_deepgemm_summary(): + logger.info( + "Enabled VLLM DeepGEMM forward+BF16-backward on %d Megatron linears (layers=%s)", + len(wrapped), + _format_int_ranges(target_layers), + ) + logger.debug("DeepGEMM wrapped Megatron linears: %s", ", ".join(wrapped)) + + +def enable_vllm_router_gemm(args, model, store_prefix: str) -> None: + """Use VLLM GateLinear forward with Megatron GEMM backward.""" + + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + del store_prefix + target_layers = _as_set( + getattr(args, "megatron_deepgemm_moe_forward_layers", None), + (), + ) + if not target_layers: + return + + wrapped = [] + for model_chunk in model: + for name, module in model_chunk.named_modules(): + if not name.endswith("mlp.router"): + continue + global_layer_index = _get_global_layer_index(model_chunk, name) + if global_layer_index not in target_layers: + continue + if getattr(module, "_vime_vllm_router_gemm_wrapped", False): + continue + if getattr(module, "bias", None) is not None: + raise RuntimeError("VLLM router alignment does not support a biased Megatron router") + + original_gating = module.gating + gate_linear = GateLinear( + input_size=module.weight.shape[1], + output_size=module.weight.shape[0], + bias=False, + out_dtype=torch.float32, + params_dtype=module.weight.dtype, + prefix="", + ) + gate_linear.weight = module.weight + object.__setattr__(module, "_vime_vllm_gate_linear", gate_linear) + + def gating( + self, + value: torch.Tensor, + ) -> torch.Tensor: + if self.weight.device.type == "cpu" and value.is_cuda: + self.weight.data = self.weight.data.to(device=value.device) + router_dtype = getattr( + getattr(self, "config", None), + "moe_router_dtype", + "fp32", + ) + if router_dtype != "fp32": + raise RuntimeError( + "VLLM persistent router alignment requires " f"moe_router_dtype=fp32, got {router_dtype}" + ) + return _VLLMRouterGEMMWithMegatronBackward.apply( + value, + self.weight, + self._vime_vllm_gate_linear, + ) + + module.gating = types.MethodType(gating, module) + module._vime_vllm_router_gemm_wrapped = True + module._vime_vllm_router_gemm_original = original_gating + wrapped.append(name) + + if wrapped and _should_log_deepgemm_summary(): + logger.info( + "Enabled VLLM batch-invariant FP32 router forward on %d routers " + "(layers=%s); Megatron backward retained", + len(wrapped), + _format_int_ranges(target_layers), + ) + + +def enable_vllm_global_batch_invariant_ops() -> None: + """Enable the same process-global deterministic operators as VLLM. + + VLLM turns these operators on from ``ModelRunner`` when deterministic + inference is requested. Megatron actors are separate processes, so + DeepGEMM's process-local batch-invariant switch alone is insufficient: + RMS reductions, BMMs, FP32 matmuls, and log-softmax would otherwise still + use Megatron's normal batch-shaped kernels. + """ + if os.environ.get("VLLM_BATCH_INVARIANT", "0").lower() not in { + "1", + "true", + "yes", + "on", + }: + return + + from vllm.model_executor.layers import batch_invariant + + batch_invariant.enable_batch_invariant_mode() + + +def _vllm_batch_invariant_rmsnorm( + value: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant + + return rms_norm_batch_invariant(value, weight, eps) + + +def _vllm_rmsnorm_from_fp32_sum( + value_fp32: torch.Tensor, + weight: torch.Tensor, + eps: float, + output_dtype: torch.dtype, +) -> torch.Tensor: + """Match VLLM's native RMSNorm on an unrounded FP32 residual sum.""" + variance = value_fp32.pow(2).mean(dim=-1, keepdim=True) + normalized = value_fp32 * torch.rsqrt(variance + eps) + return (normalized * weight).to(output_dtype) + + +def enable_vllm_layer0_input_rmsnorm( + args, + model, + store_prefix: str, +) -> None: + """Match standalone replay's independent RMSNorm at every PP boundary. + + Later Transformer layers consume the explicit FP32 residual sum and use + Megatron's VLLM-aligned fused residual/RMSNorm path. The first local + layer on each pipeline stage has no local preceding residual sum, so it + still calls its standalone TE RMSNorm module. The standalone accuracy + replay replaces all such modules with VLLM's batch-invariant RMSNorm; do + the same here. Checking for global layer zero is insufficient with PP>1: + stage 1 in the canonical PP8 layout starts at global layer 2. + + Training uses an analytic RMSNorm backward, so the original TE forward is + not evaluated a second time. + """ + del args, store_prefix + if os.environ.get("MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS", "0") != "1": + return + + if isinstance(model, torch.nn.Module): + model_chunks = [model] + else: + try: + model_chunks = list(model) + except TypeError: + return + + patched = [] + for model_chunk in model_chunks: + local_layers = [] + for layer_name, layer in model_chunk.named_modules(): + match = re.match( + r"^(?:.*\.)?decoder\.layers\.(\d+)$", + layer_name, + ) + if match is not None: + local_layers.append((int(match.group(1)), layer_name, layer)) + if not local_layers: + continue + + # Each model chunk owns one contiguous local decoder. Its first layer + # is the independent RMS boundary even when its global layer number is + # greater than one because it follows a pipeline send/recv. + _, layer_name, layer = min(local_layers, key=lambda item: item[0]) + module = getattr(layer, "input_layernorm", None) + if module is None or getattr(module, "_vime_vllm_pipeline_input_rmsnorm_wrapped", False): + continue + if not hasattr(module, "weight") or not hasattr(module, "eps"): + raise RuntimeError(f"{layer_name}.input_layernorm is missing RMSNorm weight/eps") + + original_forward = module.forward + + def forward( + patched_module: torch.nn.Module, + value: torch.Tensor, + ) -> torch.Tensor: + return _VLLMRMSNormWithAnalyticBackward.apply( + value.detach(), + value, + patched_module.weight, + float(patched_module.eps), + False, + value.dtype, + ) + + module.forward = types.MethodType(forward, module) + module._vime_vllm_pipeline_input_rmsnorm_wrapped = True + module._vime_vllm_pipeline_input_rmsnorm_original = original_forward + patched.append(f"{layer_name}.input_layernorm") + + if patched and _should_log_deepgemm_summary(): + logger.info( + "Enabled single-pass VLLM batch-invariant pipeline-stage input " + "RMSNorm with analytic backward on %d module(s)", + len(patched), + ) + + +def enable_vllm_absorbed_kv_rmsnorm( + args, + model, + store_prefix: str, +) -> None: + """Match VLLM's RMSNorm for the absorbed MLA KV latent. + + MCore invokes ``torch.nn.functional.rms_norm`` directly for this value, so + neither the standalone TE RMSNorm replacement nor the fused-linear wrapper + reaches it. The standalone accuracy harness scopes a replacement to each + attention forward. Reproduce that behavior with a direct analytic + backward instead of evaluating a native RMSNorm surrogate. + """ + del store_prefix + if os.environ.get("MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS", "0") != "1": + return + + target_layers = _as_set( + getattr(args, "megatron_deepgemm_forward_layers", None), + (), + ) + if isinstance(model, torch.nn.Module): + model_chunks = [model] + else: + try: + model_chunks = list(model) + except TypeError: + return + + wrapped = [] + for model_chunk in model_chunks: + for layer_name, layer in model_chunk.named_modules(): + if re.match(r"^(?:.*\.)?decoder\.layers\.\d+$", layer_name) is None: + continue + global_layer = int(getattr(layer, "layer_number", 0)) - 1 + if target_layers and global_layer not in target_layers: + continue + attention = getattr(layer, "self_attention", None) + kv_up = getattr(attention, "linear_kv_up_proj", None) + target_weight = getattr(kv_up, "layer_norm_weight", None) + if attention is None or target_weight is None: + continue + if getattr(attention, "_vime_vllm_absorbed_kv_rmsnorm_wrapped", False): + continue + + original_forward = attention.forward + + def forward( + patched_attention: torch.nn.Module, + *forward_args, + original_forward=original_forward, + target_weight=target_weight, + **forward_kwargs, + ): + original_rms_norm = torch.nn.functional.rms_norm + + def matched_rms_norm( + value: torch.Tensor, + normalized_shape, + weight: torch.Tensor | None = None, + eps: float | None = None, + ) -> torch.Tensor: + expected_shape = (target_weight.numel(),) + if value.shape[-1] != target_weight.numel() or tuple(normalized_shape) != expected_shape: + return original_rms_norm( + value, + normalized_shape, + weight=weight, + eps=eps, + ) + + effective_eps = float(patched_attention.config.layernorm_epsilon if eps is None else eps) + return _VLLMRMSNormWithAnalyticBackward.apply( + value.to(dtype=target_weight.dtype).detach(), + value, + target_weight, + effective_eps, + False, + target_weight.dtype, + ) + + torch.nn.functional.rms_norm = matched_rms_norm + try: + return original_forward(*forward_args, **forward_kwargs) + finally: + torch.nn.functional.rms_norm = original_rms_norm + + attention.forward = types.MethodType(forward, attention) + attention._vime_vllm_absorbed_kv_rmsnorm_wrapped = True + attention._vime_vllm_absorbed_kv_rmsnorm_original = original_forward + wrapped.append(f"{layer_name}.self_attention") + + if wrapped and _should_log_deepgemm_summary(): + logger.info( + "Enabled VLLM batch-invariant absorbed-KV RMSNorm forward " + "with analytic backward on %d attention module(s) (layers=%s)", + len(wrapped), + _format_int_ranges(target_layers), + ) + + +def enable_vllm_final_rmsnorm( + args, + model, + store_prefix: str, +) -> None: + """Match the standalone replay's final RMSNorm visible forward. + + Compact layer outputs are captured before ``decoder.final_layernorm``. + Consequently, all Transformer layers can be bitwise identical while the + final logprobs still differ if the training path keeps Transformer + Engine's final RMS kernel. The standalone replay uses VLLM's RMSNorm + and, when available, normalizes the unrounded FP32 residual sum retained + by the last Transformer layer. + + Use an analytic RMSNorm backward evaluated from the same input that the + original final norm consumed. When an exact FP32 residual is available, + it remains visible-forward-only to preserve the established gradient + topology without executing the original TE forward. + """ + del args, store_prefix + if os.environ.get("MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS", "0") != "1": + return + + if isinstance(model, torch.nn.Module): + model_chunks = [model] + else: + try: + model_chunks = list(model) + except TypeError: + return + + patched = [] + for model_chunk in model_chunks: + named_modules = dict(model_chunk.named_modules()) + for module_name, module in named_modules.items(): + if not module_name.endswith("decoder.final_layernorm"): + continue + if getattr(module, "_vime_vllm_final_rmsnorm_wrapped", False): + continue + if not hasattr(module, "weight") or not hasattr(module, "eps"): + raise RuntimeError(f"{module_name} is missing RMSNorm weight/eps") + + decoder_name = module_name.removesuffix(".final_layernorm") + decoder = named_modules.get(decoder_name) + layers = getattr(decoder, "layers", None) + residual_source = layers[-1] if layers else None + original_forward = module.forward + + def forward( + patched_module: torch.nn.Module, + value: torch.Tensor, + *, + residual_source=residual_source, + ) -> torch.Tensor: + exact_residual_sum = getattr( + value, + "_vllm_residual_sum_fp32", + None, + ) + if exact_residual_sum is None and residual_source is not None: + exact_residual_sum = getattr( + residual_source, + "_vllm_residual_sum_fp32", + None, + ) + + visible_input = value if exact_residual_sum is None else exact_residual_sum + return _VLLMRMSNormWithAnalyticBackward.apply( + visible_input.detach(), + value, + patched_module.weight, + float(patched_module.eps), + exact_residual_sum is not None, + value.dtype, + ) + + module.forward = types.MethodType(forward, module) + module._vime_vllm_final_rmsnorm_wrapped = True + module._vime_vllm_final_rmsnorm_original = original_forward + patched.append(module_name) + + if patched and _should_log_deepgemm_summary(): + logger.info( + "Enabled single-pass VLLM-aligned final RMSNorm with analytic " "backward on %d module(s)", + len(patched), + ) + + +def _vllm_swiglu_with_megatron_backward(value: torch.Tensor) -> torch.Tensor: + """Use one VLLM BF16 SwiGLU forward with an analytic backward.""" + if value.shape[-1] % 2: + raise RuntimeError(f"SwiGLU input width must be even, got {value.shape[-1]}") + return _VLLMSwiGLUWithAnalyticBackward.apply(value) + + +def _wrap_vllm_swiglu_mlp( + module: torch.nn.Module, + module_name: str, + *, + return_tuple: bool, +) -> bool: + if getattr(module, "_vime_vllm_swiglu_wrapped", False): + return False + if not hasattr(module, "linear_fc1") or not hasattr(module, "linear_fc2"): + raise RuntimeError(f"SwiGLU target {module_name} is not an MLP") + + def forward( + self, + hidden_states: torch.Tensor, + per_token_scale: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + ): + # Current Megatron passes this keyword to both dense MLPs and MoE + # layers. Its native dense MLP accepts and ignores it; padding tokens + # are excluded by the loss mask rather than altered inside the MLP. + del padding_mask + if per_token_scale is not None: + raise RuntimeError(f"VLLM SwiGLU alignment does not support per_token_scale: {module_name}") + gate_up, bias = self.linear_fc1(hidden_states) + if bias is not None: + raise RuntimeError(f"VLLM SwiGLU alignment requires a bias-free MLP: {module_name}") + down_input = _vllm_swiglu_with_megatron_backward(gate_up) + output, output_bias = self.linear_fc2(down_input) + if output_bias is not None: + raise RuntimeError(f"VLLM SwiGLU alignment requires a bias-free MLP: {module_name}") + + if not return_tuple: + if getattr(self, "use_shared_expert_gate", False): + logits = torch.nn.functional.linear(hidden_states, self.gate_weight) + output = output * torch.sigmoid(logits) + return output + return output, None + + module.forward = types.MethodType(forward, module) + module._vime_vllm_swiglu_wrapped = True + module._vime_vllm_swiglu_module_name = module_name + return True + + +def enable_vllm_swiglu_forward(args, model, store_prefix: str) -> None: + """Match VLLM's fused SwiGLU on dense and shared-expert MLPs.""" + del store_prefix + target_layers = _as_set( + getattr(args, "megatron_deepgemm_forward_layers", None), + (), + ) + if not target_layers: + return + + if isinstance(model, torch.nn.Module): + model_chunks = [model] + else: + try: + model_chunks = list(model) + except TypeError: + return + + wrapped = [] + for model_chunk in model_chunks: + for layer_name, layer in model_chunk.named_modules(): + match = re.match(r"^(?:.*\.)?decoder\.layers\.\d+$", layer_name) + if match is None: + continue + layer_number = getattr(layer, "layer_number", None) + if layer_number is None or int(layer_number) - 1 not in target_layers: + continue + mlp = getattr(layer, "mlp", None) + if mlp is None: + continue + if hasattr(mlp, "experts"): + shared = getattr(mlp, "shared_experts", None) + if shared is not None and _wrap_vllm_swiglu_mlp( + shared, + f"{layer_name}.mlp.shared_experts", + return_tuple=False, + ): + wrapped.append(f"{layer_name}.mlp.shared_experts") + elif _wrap_vllm_swiglu_mlp( + mlp, + f"{layer_name}.mlp", + return_tuple=True, + ): + wrapped.append(f"{layer_name}.mlp") + + if wrapped and _should_log_deepgemm_summary(): + logger.info( + "Enabled single-pass VLLM SwiGLU forward with analytic backward on %d MLPs " "(layers=%s)", + len(wrapped), + _format_int_ranges(target_layers), + ) + + +def _enable_deepgemm_all_forward( + args, + model, + store_prefix: str, +) -> None: + """Install selected dense/indexer/shared and routed-MoE forward paths.""" + linear_layers = getattr(args, "megatron_deepgemm_forward_layers", None) + moe_layers = getattr(args, "megatron_deepgemm_moe_forward_layers", None) + if not linear_layers and not moe_layers: + raise RuntimeError( + "The combined DeepGEMM hook requires at least one of " + "--megatron-deepgemm-forward-layers or " + "--megatron-deepgemm-moe-forward-layers" + ) + + enable_vllm_global_batch_invariant_ops() + enable_vllm_layer0_input_rmsnorm(args, model, store_prefix) + enable_vllm_absorbed_kv_rmsnorm(args, model, store_prefix) + enable_vllm_final_rmsnorm(args, model, store_prefix) + + if linear_layers: + enable_deepgemm_forward(args, model, store_prefix) + enable_vllm_swiglu_forward(args, model, store_prefix) + if moe_layers: + from vime.backends.megatron_utils.alignment.deepgemm_moe_forward import ( + enable_deepgemm_moe_forward, + enable_vllm_deepep_moe_alignment, + ) + + enable_deepgemm_moe_forward(args, model, store_prefix) + if os.environ.get("MEGATRON_USE_VLLM_ROUTER_GEMM", "0") == "1": + enable_vllm_router_gemm(args, model, store_prefix) + enable_vllm_deepep_moe_alignment(args, model, store_prefix) + + if os.getenv("VIME_LAYERWISE_ALIGNMENT_DUMP_DIR"): + from vime.backends.megatron_utils.alignment.layerwise_alignment import enable_megatron_layerwise_dump + + enable_megatron_layerwise_dump(args, model, store_prefix) + + +def enable_deepgemm_all_forward(args, model, store_prefix: str) -> None: + """Install every selected alignment path (native LM head).""" + _enable_deepgemm_all_forward(args, model, store_prefix) + + +def enable_deepgemm_all_forward_before_train_step( + args, + rollout_id: int, + step_id: int, + model, + optimizer, + opt_param_scheduler, +) -> None: + """Install alignment with the native LM head before the train step.""" + del rollout_id, step_id, opt_param_scheduler + enable_deepgemm_all_forward(args, model, store_prefix="") + del optimizer diff --git a/vime/backends/megatron_utils/alignment/deepgemm_moe_forward.py b/vime/backends/megatron_utils/alignment/deepgemm_moe_forward.py new file mode 100644 index 000000000..bc3922a68 --- /dev/null +++ b/vime/backends/megatron_utils/alignment/deepgemm_moe_forward.py @@ -0,0 +1,3164 @@ +"""Align Megatron TEGroupedMLP with VLLM's DeepGEMM MoE path. + + block-FP8 fc1 -> SwiGLU -> block-FP8 fc2 -> FP32 router-probability multiply + +Moving the probability multiply after fc2 matches VLLM more closely than +wrapping the two grouped linears independently. The wrapper installs a custom +autograd function: + + forward = block-FP8 grouped DeepGEMM; + backward = grouped BF16 dgrad/wgrad GEMMs plus analytic + SwiGLU/router-probability gradients. +""" + +from __future__ import annotations + +import logging +import math +import os +import re +import types +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from megatron.core import parallel_state + +from vime.backends.megatron_utils.alignment.deepgemm_forward import ( + _deepgemm_bf16_gemm_nn, + _deepgemm_bf16_gemm_nt, + _deepgemm_bf16_gemm_tn, + _format_int_ranges, + _should_log_deepgemm_summary, + _sum_to_parameter_dtype, + _vllm_silu_and_mul, +) + +logger = logging.getLogger(__name__) + +_BLOCK_SIZE = 128 +_GROUPED_M_ALIGNMENT = 128 +_ROUTER_PROBABILITY_CHUNK_ROWS = 1024 +_UNPAD_CHUNK_ROWS = 1024 +_BACKWARD_CHUNK_ROWS = 1024 +_SWIGLU_POINTWISE_CHUNK_ROWS = 512 +_DEFAULT_EXPERTS_PER_GROUP = 4 +_DEFAULT_BACKWARD_EXPERTS_PER_GROUP = 4 +_DEFAULT_BACKWARD_MAX_PADDED_BYTES = 256 * 1024 * 1024 +_DEFAULT_TARGET_SUFFIXES = ("mlp.experts",) +_PREALLOCATED_COMBINE_BUFFER_ATTR = "_vime_preallocated_combine_buffer" +_PREALLOCATED_TOKEN_COMBINE_ATTR = "_vime_preallocated_token_combine" +_COMBINE_WORKSPACE_ATTR = "_vime_combine_workspace" +_LAYER_PATH_RE = re.compile(r"^(?P(?:.*\.)?decoder\.layers\.(?P\d+))(?:\.|$)") + + +@dataclass(frozen=True) +class _MoELayout: + num_local_experts: int + hidden_size: int + ffn_hidden_size: int + + @property + def fc1_weight_shape(self) -> tuple[int, int]: + return (2 * self.ffn_hidden_size, self.hidden_size) + + @property + def fc2_weight_shape(self) -> tuple[int, int]: + return (self.hidden_size, self.ffn_hidden_size) + + +@dataclass(frozen=True) +class _DeepGEMMOps: + quantize_weight: Callable[..., tuple[torch.Tensor, torch.Tensor]] + quantize_activation: Callable[..., tuple[torch.Tensor, torch.Tensor]] + align_input_scale: Callable[[torch.Tensor], torch.Tensor] + grouped_gemm: Callable[..., Any] + silu_and_mul: Callable[[torch.Tensor, torch.Tensor], Any] + # Blackwell (sm100+) uses UE8M0 (power-of-two) block scales; Hopper (sm90) + # uses FP32 block scales. When ``scale_ue8m0`` is False the H100 path below + # is byte-for-byte unchanged. + scale_ue8m0: bool = False + need_tma_aligned_scales: bool = True + transform_weight_scale: Callable[..., torch.Tensor] | None = None + + +def _apply_router_probability_fp32_inplace( + down_output: torch.Tensor, + permuted_probs: torch.Tensor, +) -> torch.Tensor: + """Apply the post-fc2 router probability without a full-size FP32 temporary. + + VLLM performs this multiply in FP32 and casts the result back to the + activation dtype. A single expression such as + ``(down_output.float() * probs.float()).to(dtype)`` temporarily materializes + the entire MoE output in FP32. For a packed 4x4 rollout this can exceed + 13 GiB per rank. Processing independent row chunks is numerically identical + while keeping the FP32 workspace bounded. + + The DeepGEMM output is scratch storage owned by this forward, so updating it + in place also avoids allocating a second full-size BF16 tensor. + """ + if down_output.ndim != 2: + raise RuntimeError(f"Expected a 2D MoE output, got shape {tuple(down_output.shape)}") + probabilities_fp32 = permuted_probs.detach().reshape(-1, 1).to(torch.float32) + if probabilities_fp32.shape[0] != down_output.shape[0]: + raise RuntimeError( + "MoE output/router probability row mismatch: " f"{down_output.shape[0]} != {probabilities_fp32.shape[0]}" + ) + + for start in range(0, down_output.shape[0], _ROUTER_PROBABILITY_CHUNK_ROWS): + end = min(start + _ROUTER_PROBABILITY_CHUNK_ROWS, down_output.shape[0]) + scaled = (down_output[start:end].to(torch.float32) * probabilities_fp32[start:end]).to(down_output.dtype) + down_output[start:end].copy_(scaled) + return down_output + + +def _router_probability_grad_fp32_chunked( + grad_output: torch.Tensor, + down_output: torch.Tensor, +) -> torch.Tensor: + """Compute the per-row router-probability gradient with bounded scratch. + + Each output row is an independent hidden-dimension dot product. Row + chunking therefore preserves the exact FP32 reduction performed by the + unchunked expression while avoiding two route-sized FP32 casts plus their + product being live at once. + """ + if grad_output.ndim != 2 or down_output.ndim != 2: + raise RuntimeError( + "Expected 2D router-gradient inputs, got " f"{tuple(grad_output.shape)} and {tuple(down_output.shape)}" + ) + if grad_output.shape != down_output.shape: + raise RuntimeError( + "Router-gradient input shape mismatch: " f"{tuple(grad_output.shape)} != {tuple(down_output.shape)}" + ) + + result = torch.empty( + (grad_output.shape[0], 1), + dtype=torch.float32, + device=grad_output.device, + ) + for start in range(0, grad_output.shape[0], _ROUTER_PROBABILITY_CHUNK_ROWS): + end = min(start + _ROUTER_PROBABILITY_CHUNK_ROWS, grad_output.shape[0]) + result[start:end].copy_( + (grad_output[start:end].float() * down_output[start:end].float()).sum( + dim=-1, + keepdim=True, + ) + ) + return result + + +def _compact_valid_rows_inplace( + padded_value: torch.Tensor, + valid_rows: torch.Tensor, +) -> torch.Tensor: + """Remove per-expert padding without a second full-size activation tensor. + + ``valid_rows`` is produced in ascending expert/row order and always + satisfies ``valid_rows[i] >= i``. Therefore copying ascending chunks into + the prefix cannot overwrite a source needed by a later chunk. Each + ``index_select`` only materializes a bounded temporary, and the returned + prefix view keeps ownership of the original DeepGEMM output storage. + """ + if padded_value.ndim != 2 or valid_rows.ndim != 1: + raise RuntimeError( + "Expected a 2D padded value and 1D valid-row indices, got " + f"{tuple(padded_value.shape)} and {tuple(valid_rows.shape)}" + ) + if valid_rows.numel() > padded_value.shape[0]: + raise RuntimeError(f"Valid-row count {valid_rows.numel()} exceeds padded rows {padded_value.shape[0]}") + if valid_rows.numel() == 0: + return padded_value.narrow(0, 0, 0) + + # Avoid adding device synchronizations to every MoE layer. CUDA indices + # come exclusively from _pad_expert_rows, which guarantees these + # invariants by construction; retain the defensive validation for CPU + # callers and unit tests. + if not valid_rows.is_cuda: + expected_positions = torch.arange( + valid_rows.numel(), + device=valid_rows.device, + dtype=valid_rows.dtype, + ) + if bool(torch.any(valid_rows < expected_positions)): + raise RuntimeError("Valid-row indices cannot move a row backward into a future source") + if bool(torch.any(valid_rows[1:] <= valid_rows[:-1])): + raise RuntimeError("Valid-row indices must be strictly increasing") + if int(valid_rows[-1].item()) >= padded_value.shape[0]: + raise RuntimeError( + f"Valid-row index {int(valid_rows[-1].item())} exceeds padded rows {padded_value.shape[0]}" + ) + + for start in range(0, valid_rows.numel(), _UNPAD_CHUNK_ROWS): + end = min(start + _UNPAD_CHUNK_ROWS, valid_rows.numel()) + selected = padded_value.index_select(0, valid_rows[start:end]) + padded_value[start:end].copy_(selected) + return padded_value.narrow(0, 0, valid_rows.numel()) + + +def _swiglu_forward_chunked( + gate: torch.Tensor, + up: torch.Tensor, +) -> torch.Tensor: + """Evaluate FP32 SwiGLU into BF16 storage with bounded temporaries.""" + if gate.shape != up.shape or gate.dtype != up.dtype: + raise RuntimeError( + "SwiGLU gate/up mismatch: " f"{tuple(gate.shape)}/{gate.dtype} != {tuple(up.shape)}/{up.dtype}" + ) + down_input = torch.empty_like(gate) + for start in range(0, gate.shape[0], _SWIGLU_POINTWISE_CHUNK_ROWS): + end = min(start + _SWIGLU_POINTWISE_CHUNK_ROWS, gate.shape[0]) + gate_f = gate[start:end].float() + up_f = up[start:end].float() + silu_gate = F.silu(gate_f) + down_input[start:end].copy_((silu_gate * up_f).to(dtype=gate.dtype)) + return down_input + + +def _swiglu_backward_chunked( + gate: torch.Tensor, + up: torch.Tensor, + grad_down_input: torch.Tensor, +) -> torch.Tensor: + """Evaluate the established FP32 SwiGLU derivative with bounded temporaries.""" + if gate.shape != up.shape or gate.shape != grad_down_input.shape: + raise RuntimeError( + "SwiGLU backward shape mismatch: " + f"{tuple(gate.shape)}, {tuple(up.shape)}, {tuple(grad_down_input.shape)}" + ) + if gate.dtype != up.dtype or gate.dtype != grad_down_input.dtype: + raise RuntimeError("SwiGLU backward dtype mismatch: " f"{gate.dtype}, {up.dtype}, {grad_down_input.dtype}") + + grad_gate_up = torch.empty( + (gate.shape[0], 2 * gate.shape[1]), + device=gate.device, + dtype=gate.dtype, + ) + grad_gate_out, grad_up_out = grad_gate_up.chunk(2, dim=-1) + for start in range(0, gate.shape[0], _SWIGLU_POINTWISE_CHUNK_ROWS): + end = min(start + _SWIGLU_POINTWISE_CHUNK_ROWS, gate.shape[0]) + gate_f = gate[start:end].float() + up_f = up[start:end].float() + grad_down_input_f = grad_down_input[start:end].float() + silu_gate = F.silu(gate_f) + sigmoid_gate = torch.sigmoid(gate_f) + grad_gate = grad_down_input_f * up_f * sigmoid_gate * (1.0 + gate_f * (1.0 - sigmoid_gate)) + grad_up = grad_down_input_f * silu_gate + grad_gate_out[start:end].copy_(grad_gate.to(dtype=gate.dtype)) + grad_up_out[start:end].copy_(grad_up.to(dtype=gate.dtype)) + return grad_gate_up + + +def _grouped_bf16_backward_experts_per_group(num_local_experts: int) -> int: + configured = os.environ.get("VIME_DEEPGEMM_MOE_BF16_BACKWARD_EXPERTS_PER_GROUP") + if configured is None: + return min(_DEFAULT_BACKWARD_EXPERTS_PER_GROUP, num_local_experts) + try: + experts_per_group = int(configured) + except ValueError as exc: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_BF16_BACKWARD_EXPERTS_PER_GROUP must be a " f"positive integer, got {configured!r}" + ) from exc + if experts_per_group <= 0: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_BF16_BACKWARD_EXPERTS_PER_GROUP must be a " f"positive integer, got {experts_per_group}" + ) + return min(experts_per_group, num_local_experts) + + +def _grouped_bf16_backward_max_padded_bytes() -> int: + configured = os.environ.get("VIME_DEEPGEMM_MOE_BF16_BACKWARD_MAX_PADDED_BYTES") + if configured is None: + return _DEFAULT_BACKWARD_MAX_PADDED_BYTES + try: + max_padded_bytes = int(configured) + except ValueError as exc: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_BF16_BACKWARD_MAX_PADDED_BYTES must be a " f"positive integer, got {configured!r}" + ) from exc + if max_padded_bytes <= 0: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_BF16_BACKWARD_MAX_PADDED_BYTES must be a " f"positive integer, got {max_padded_bytes}" + ) + return max_padded_bytes + + +def _padded_expert_hidden_bytes( + count: int, + *, + hidden_size: int, + element_size: int, +) -> int: + padded_count = ((count + _GROUPED_M_ALIGNMENT - 1) // _GROUPED_M_ALIGNMENT) * _GROUPED_M_ALIGNMENT if count else 0 + return padded_count * hidden_size * element_size + + +def _grouped_bf16_backward_expert_ranges( + counts: tuple[int, ...], + *, + hidden_size: int, + element_size: int, +) -> tuple[tuple[int, int], ...]: + """Greedily group adjacent experts without exceeding the padded-input cap.""" + experts_per_group = _grouped_bf16_backward_experts_per_group(len(counts)) + max_padded_bytes = _grouped_bf16_backward_max_padded_bytes() + ranges = [] + expert_start = 0 + while expert_start < len(counts): + expert_end = expert_start + group_padded_bytes = 0 + while expert_end < len(counts) and expert_end - expert_start < experts_per_group: + expert_padded_bytes = _padded_expert_hidden_bytes( + counts[expert_end], + hidden_size=hidden_size, + element_size=element_size, + ) + if expert_end > expert_start and group_padded_bytes + expert_padded_bytes > max_padded_bytes: + break + group_padded_bytes += expert_padded_bytes + expert_end += 1 + ranges.append((expert_start, expert_end)) + expert_start = expert_end + return tuple(ranges) + + +def _use_grouped_bf16_backward( + hidden_states: torch.Tensor, + counts: tuple[int, ...], + needs_fc1_weights: tuple[bool, ...], + needs_fc2_weights: tuple[bool, ...], +) -> bool: + enabled = os.environ.get("VIME_DEEPGEMM_MOE_GROUPED_BF16_BACKWARD", "1").lower() in { + "1", + "true", + "yes", + "on", + } + max_padded_bytes = _grouped_bf16_backward_max_padded_bytes() + largest_expert_bytes = max( + ( + _padded_expert_hidden_bytes( + count, + hidden_size=hidden_states.shape[1], + element_size=hidden_states.element_size(), + ) + for count in counts + ), + default=0, + ) + if ( + not enabled + or not hidden_states.is_cuda + or hidden_states.dtype != torch.bfloat16 + # A single expert cannot be split by the contiguous grouped kernel. + # Fall back to the established 1024-row path for that rare hot-expert + # case instead of risking a full-padded allocation OOM. + or largest_expert_bytes > max_padded_bytes + ): + return False + + if any(needs_fc1_weights) or any(needs_fc2_weights): + import deep_gemm + + if not hasattr(deep_gemm, "k_grouped_bf16_gemm_tn_contiguous"): + return False + return True + + +def _deepgemm_bf16_m_grouped_gemm_nt( + lhs: torch.Tensor, + rhs: torch.Tensor, + grouped_layout: torch.Tensor, +) -> torch.Tensor: + """Compute contiguous expert-major ``lhs @ rhs[group].T`` in one launch.""" + if lhs.ndim != 2 or rhs.ndim != 3 or lhs.shape[1] != rhs.shape[2]: + raise RuntimeError(f"Grouped BF16 NT GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if grouped_layout.shape != (lhs.shape[0],) or grouped_layout.dtype != torch.int32: + raise RuntimeError( + "Grouped BF16 NT layout mismatch: " + f"{tuple(grouped_layout.shape)}/{grouped_layout.dtype} for {lhs.shape[0]} rows" + ) + if not (lhs.is_cuda and rhs.is_cuda and grouped_layout.is_cuda): + raise RuntimeError("Grouped BF16 backward requires CUDA tensors") + import deep_gemm + + out = torch.empty((lhs.shape[0], rhs.shape[1]), device=lhs.device, dtype=torch.bfloat16) + deep_gemm.m_grouped_bf16_gemm_nt_contiguous( + lhs.contiguous(), + rhs.contiguous(), + out, + grouped_layout.contiguous(), + ) + return out + + +def _deepgemm_bf16_m_grouped_gemm_nn( + lhs: torch.Tensor, + rhs: torch.Tensor, + grouped_layout: torch.Tensor, +) -> torch.Tensor: + """Compute contiguous expert-major ``lhs @ rhs[group]`` in one launch.""" + if lhs.ndim != 2 or rhs.ndim != 3 or lhs.shape[1] != rhs.shape[1]: + raise RuntimeError(f"Grouped BF16 NN GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if grouped_layout.shape != (lhs.shape[0],) or grouped_layout.dtype != torch.int32: + raise RuntimeError( + "Grouped BF16 NN layout mismatch: " + f"{tuple(grouped_layout.shape)}/{grouped_layout.dtype} for {lhs.shape[0]} rows" + ) + if not (lhs.is_cuda and rhs.is_cuda and grouped_layout.is_cuda): + raise RuntimeError("Grouped BF16 backward requires CUDA tensors") + import deep_gemm + + out = torch.empty((lhs.shape[0], rhs.shape[2]), device=lhs.device, dtype=torch.bfloat16) + deep_gemm.m_grouped_bf16_gemm_nn_contiguous( + lhs.contiguous(), + rhs.contiguous(), + out, + grouped_layout.contiguous(), + ) + return out + + +def _deepgemm_bf16_k_grouped_gemm_tn( + lhs: torch.Tensor, + rhs: torch.Tensor, + grouped_k: tuple[int, ...], +) -> torch.Tensor: + """Compute per-expert ``lhs.T @ rhs`` weight gradients in one launch.""" + if lhs.ndim != 2 or rhs.ndim != 2 or lhs.shape[0] != rhs.shape[0]: + raise RuntimeError(f"Grouped BF16 TN GEMM shape mismatch: {tuple(lhs.shape)} x {tuple(rhs.shape)}") + if sum(grouped_k) != lhs.shape[0]: + raise RuntimeError( + "Grouped BF16 TN row-count mismatch: " f"sum({grouped_k})={sum(grouped_k)} != {lhs.shape[0]}" + ) + if not grouped_k: + raise RuntimeError("Grouped BF16 TN GEMM requires at least one expert") + if any(k < 0 or k % _GROUPED_M_ALIGNMENT for k in grouped_k): + raise RuntimeError( + f"Grouped BF16 TN K sizes must be non-negative multiples of {_GROUPED_M_ALIGNMENT}: {grouped_k}" + ) + if not (lhs.is_cuda and rhs.is_cuda): + raise RuntimeError("Grouped BF16 TN GEMM requires CUDA tensors") + if lhs.dtype != torch.bfloat16 or rhs.dtype != torch.bfloat16: + raise RuntimeError(f"Grouped BF16 TN GEMM requires BF16 tensors, got {lhs.dtype} and {rhs.dtype}") + + import deep_gemm + + # The final Megatron expert gradients are BF16. Asking the grouped kernel + # to write BF16 directly matches one full per-expert DeepGEMM TN launch and + # avoids retaining an additional FP32 copy of every expert weight gradient. + out = torch.zeros( + (len(grouped_k), lhs.shape[1], rhs.shape[1]), + device=lhs.device, + dtype=torch.bfloat16, + ) + grouped_k_tensor = torch.tensor(grouped_k, device=lhs.device, dtype=torch.int32) + deep_gemm.k_grouped_bf16_gemm_tn_contiguous( + lhs.contiguous(), + rhs.contiguous(), + out, + grouped_k, + grouped_k_tensor, + out, + ) + return out + + +def _grouped_expert_backward( + *, + hidden_states: torch.Tensor, + permuted_probs: torch.Tensor, + grad_output: torch.Tensor, + fc1_weights: tuple[torch.Tensor, ...], + fc2_weights: tuple[torch.Tensor, ...], + counts: tuple[int, ...], + layout: _MoELayout, + needs_hidden: bool, + needs_probs: bool, + needs_fc1_weights: tuple[bool, ...], + needs_fc2_weights: tuple[bool, ...], + defer_router_probabilities: bool, + grad_hidden: torch.Tensor | None, + grad_probs: torch.Tensor | None, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + list[torch.Tensor | None], + list[torch.Tensor | None], +]: + """Grouped-BF16 dgrad/wgrad without changing the aligned forward.""" + probabilities = permuted_probs.reshape(-1, 1) + token_offset = 0 + grad_fc1_weights: list[torch.Tensor | None] = [None] * layout.num_local_experts + grad_fc2_weights: list[torch.Tensor | None] = [None] * layout.num_local_experts + + expert_ranges = _grouped_bf16_backward_expert_ranges( + counts, + hidden_size=hidden_states.shape[1], + element_size=hidden_states.element_size(), + ) + for expert_start, expert_end in expert_ranges: + group_counts = counts[expert_start:expert_end] + group_tokens = sum(group_counts) + if group_tokens == 0: + for expert_index in range(expert_start, expert_end): + if needs_fc1_weights[expert_index]: + grad_fc1_weights[expert_index] = torch.zeros_like(fc1_weights[expert_index]) + if needs_fc2_weights[expert_index]: + grad_fc2_weights[expert_index] = torch.zeros_like(fc2_weights[expert_index]) + continue + + hidden = hidden_states.narrow(0, token_offset, group_tokens) + grad = grad_output.narrow(0, token_offset, group_tokens) + probability = probabilities.narrow(0, token_offset, group_tokens) + padded_hidden, padded_counts, valid_rows = _pad_expert_rows(hidden, group_counts) + + grouped_layout = _build_m_indices(padded_counts, device=hidden_states.device) + fc1_group = torch.stack(fc1_weights[expert_start:expert_end], dim=0) + gate_up = _deepgemm_bf16_m_grouped_gemm_nt( + padded_hidden, + fc1_group, + grouped_layout, + ) + needs_fc1_group = needs_fc1_weights[expert_start:expert_end] + needs_fc2_group = needs_fc2_weights[expert_start:expert_end] + if not any(needs_fc1_group): + # The FC1 input is otherwise not needed again in backward. + del padded_hidden + gate, up = gate_up.chunk(2, dim=-1) + + fc2_group = torch.stack(fc2_weights[expert_start:expert_end], dim=0) + padded_grad, grad_padded_counts, grad_valid_rows = _pad_expert_rows(grad, group_counts) + if grad_padded_counts != padded_counts or (valid_rows is None) != (grad_valid_rows is None): + raise RuntimeError("Grouped BF16 backward produced inconsistent gradient padding") + + padded_probability = None + probability_valid_rows = None + if not defer_router_probabilities: + padded_probability, probability_padded_counts, probability_valid_rows = _pad_expert_rows( + probability, + group_counts, + ) + if probability_padded_counts != padded_counts or (valid_rows is None) != (probability_valid_rows is None): + raise RuntimeError("Grouped BF16 backward produced inconsistent probability padding") + + down_input = None + if (needs_probs and not defer_router_probabilities) or any(needs_fc2_group): + down_input = _swiglu_forward_chunked(gate, up) + if needs_probs and not defer_router_probabilities: + assert down_input is not None + down_output = _deepgemm_bf16_m_grouped_gemm_nt( + down_input, + fc2_group, + grouped_layout, + ) + grad_probs_padded = (padded_grad.float() * down_output.float()).sum(dim=-1, keepdim=True) + if valid_rows is not None: + grad_probs_group = _compact_valid_rows_inplace(grad_probs_padded, valid_rows) + else: + grad_probs_group = grad_probs_padded + assert grad_probs is not None + grad_probs.narrow(0, token_offset, group_tokens).copy_( + grad_probs_group.reshape_as(permuted_probs.narrow(0, token_offset, group_tokens)).to( + dtype=permuted_probs.dtype + ) + ) + del down_output, grad_probs_padded, grad_probs_group + + if defer_router_probabilities: + grad_down_output = padded_grad + else: + assert padded_probability is not None + grad_down_output = (padded_grad.float() * padded_probability.float()).to(dtype=hidden_states.dtype) + if any(needs_fc2_group): + assert down_input is not None + grad_fc2_group = _deepgemm_bf16_k_grouped_gemm_tn( + grad_down_output, + down_input, + padded_counts, + ) + for local_expert_index, needs_weight in enumerate(needs_fc2_group): + if needs_weight: + grad_fc2_weights[expert_start + local_expert_index] = grad_fc2_group[local_expert_index] + del grad_fc2_group + del down_input + del padded_probability + grad_down_input = _deepgemm_bf16_m_grouped_gemm_nn( + grad_down_output, + fc2_group, + grouped_layout, + ) + del grad_down_output, padded_grad, fc2_group + + grad_gate_up = _swiglu_backward_chunked(gate, up, grad_down_input) + del grad_down_input, gate, up, gate_up + + if any(needs_fc1_group): + grad_fc1_group = _deepgemm_bf16_k_grouped_gemm_tn( + grad_gate_up, + padded_hidden, + padded_counts, + ) + for local_expert_index, needs_weight in enumerate(needs_fc1_group): + if needs_weight: + grad_fc1_weights[expert_start + local_expert_index] = grad_fc1_group[local_expert_index] + del grad_fc1_group, padded_hidden + + if needs_hidden: + grad_hidden_padded = _deepgemm_bf16_m_grouped_gemm_nn( + grad_gate_up, + fc1_group, + grouped_layout, + ) + if valid_rows is not None: + grad_hidden_group = _compact_valid_rows_inplace(grad_hidden_padded, valid_rows) + else: + grad_hidden_group = grad_hidden_padded + assert grad_hidden is not None + grad_hidden.narrow(0, token_offset, group_tokens).copy_(grad_hidden_group) + del grad_hidden_padded, grad_hidden_group + + token_offset += group_tokens + del ( + valid_rows, + grad_valid_rows, + probability_valid_rows, + grouped_layout, + fc1_group, + grad_gate_up, + ) + + return ( + grad_hidden, + None if defer_router_probabilities else grad_probs, + grad_fc1_weights, + grad_fc2_weights, + ) + + +def _ordered_route_backward( + *, + route_values: torch.Tensor, + topk_weights: torch.Tensor, + output_index: torch.Tensor, + grad_output: torch.Tensor, + grad_routes: torch.Tensor | None, + grad_weights: torch.Tensor | None, + static_mapping_valid: bool | None = None, +) -> None: + """Differentiate the ordered top-k gather with an optional static fast path.""" + routes_alias_values = ( + grad_routes is not None + and grad_routes.untyped_storage().data_ptr() == route_values.untyped_storage().data_ptr() + ) + use_static_mapping = static_mapping_valid is not None + if use_static_mapping and static_mapping_valid is None: + # The hot DeepEP caller passes this bit from its forward scatter, + # where torch.nonzero has already exposed the route count to Python. + # Other internal callers retain a safe fallback instead of assuming + # that padded token slots own an expert-output row. + static_mapping_valid = bool(torch.all(output_index >= 0).item()) + if use_static_mapping and static_mapping_valid: + # Dropless fixed-top-k routing produces exactly one valid route row for + # every flattened [token, top-k] slot. Express that static mapping + # directly instead of materializing torch.nonzero's data-dependent + # output and synchronizing once per MoE layer. + num_routes = output_index.numel() + topk = output_index.shape[1] + flat_output_index = output_index.reshape(-1) + flat_weights = topk_weights.reshape(-1) + flat_grad_weights = grad_weights.reshape(-1) if grad_weights is not None else None + chunk_rows = _BACKWARD_CHUNK_ROWS + # When the forward combine input is dead after this operation, its + # route-sized storage can hold the route gradient. Complete every + # probability gradient first because it still reads the original + # route values; only then overwrite that storage with route gradients. + if routes_alias_values and flat_grad_weights is not None: + for start in range(0, num_routes, chunk_rows): + end = min(start + chunk_rows, num_routes) + flat_positions = torch.arange( + start, + end, + device=output_index.device, + dtype=torch.long, + ) + token_rows = torch.div(flat_positions, topk, rounding_mode="floor") + route_rows = flat_output_index[start:end].to(dtype=torch.long) + token_grads = grad_output.index_select(0, token_rows) + selected_route_values = route_values.index_select(0, route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + flat_grad_weights[start:end].copy_(weight_grads) + + fused_route_grad = False + if grad_routes is not None and grad_output.is_cuda: + from vime.backends.megatron_utils.alignment.deterministic_route_kernels import ordered_route_grad + + ordered_route_grad( + grad_output.contiguous(), + topk_weights.contiguous(), + output_index.contiguous(), + grad_routes, + ) + fused_route_grad = True + + if (grad_routes is not None and not fused_route_grad) or ( + flat_grad_weights is not None and not routes_alias_values + ): + for start in range(0, num_routes, chunk_rows): + end = min(start + chunk_rows, num_routes) + flat_positions = torch.arange( + start, + end, + device=output_index.device, + dtype=torch.long, + ) + token_rows = torch.div(flat_positions, topk, rounding_mode="floor") + route_rows = flat_output_index[start:end].to(dtype=torch.long) + token_grads = grad_output.index_select(0, token_rows) + weights = flat_weights[start:end] + + if grad_routes is not None and not fused_route_grad: + route_grads = (token_grads.float() * weights.float().unsqueeze(1)).to(dtype=grad_routes.dtype) + # Every static top-k slot maps to one distinct route row. + grad_routes.index_copy_(0, route_rows, route_grads) + if flat_grad_weights is not None and not routes_alias_values: + selected_route_values = route_values.index_select(0, route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + flat_grad_weights[start:end].copy_(weight_grads) + return + + if use_static_mapping: + # Padding-aware fixed-shape path. Mapping every masked slot to row 0 + # keeps all intermediates bounded by chunk_rows and avoids allocating + # torch.nonzero's data-dependent route table. Masked slots write an + # exact zero; row 0 is restored in a final ordered kernel, so duplicate + # sentinel writes cannot affect the visible gradient. + num_slots = output_index.numel() + num_route_rows = route_values.shape[0] + flat_output_index = output_index.reshape(-1) + flat_weights = topk_weights.reshape(-1) + flat_grad_weights = grad_weights.reshape(-1) if grad_weights is not None else None + chunk_rows = _BACKWARD_CHUNK_ROWS + + if num_route_rows == 0: + if grad_routes is not None: + grad_routes.zero_() + if flat_grad_weights is not None: + flat_grad_weights.zero_() + return + + def probability_grad_chunk(start: int, end: int) -> None: + flat_positions = torch.arange( + start, + end, + device=output_index.device, + dtype=torch.long, + ) + token_rows = torch.div( + flat_positions, + output_index.shape[1], + rounding_mode="floor", + ) + route_rows = flat_output_index[start:end].to(dtype=torch.long) + valid_rows = route_rows >= 0 + safe_route_rows = route_rows.clamp(min=0, max=num_route_rows - 1) + token_grads = grad_output.index_select(0, token_rows) + selected_route_values = route_values.index_select(0, safe_route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + weight_grads.masked_fill_(~valid_rows, 0) + flat_grad_weights[start:end].copy_(weight_grads) + + if routes_alias_values and flat_grad_weights is not None: + for start in range(0, num_slots, chunk_rows): + probability_grad_chunk(start, min(start + chunk_rows, num_slots)) + + if grad_routes is not None and routes_alias_values: + grad_routes.zero_() + + for start in range(0, num_slots, chunk_rows): + end = min(start + chunk_rows, num_slots) + flat_positions = torch.arange( + start, + end, + device=output_index.device, + dtype=torch.long, + ) + token_rows = torch.div( + flat_positions, + output_index.shape[1], + rounding_mode="floor", + ) + route_rows = flat_output_index[start:end].to(dtype=torch.long) + valid_rows = route_rows >= 0 + safe_route_rows = route_rows.clamp(min=0, max=num_route_rows - 1) + token_grads = grad_output.index_select(0, token_rows) + + if grad_routes is not None: + weights = flat_weights[start:end].float().masked_fill(~valid_rows, 0) + route_grads = (token_grads.float() * weights.unsqueeze(1)).to(dtype=grad_routes.dtype) + grad_routes.index_copy_(0, safe_route_rows, route_grads) + if flat_grad_weights is not None and not routes_alias_values: + selected_route_values = route_values.index_select(0, safe_route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + weight_grads.masked_fill_(~valid_rows, 0) + flat_grad_weights[start:end].copy_(weight_grads) + + if grad_routes is not None: + route_zero_matches = flat_output_index == 0 + torch._assert_async( + torch.any(route_zero_matches), + "non-empty expert output has no route mapped to row zero", + ) + route_zero_position = torch.argmax(route_zero_matches.to(dtype=torch.int32)).reshape(1) + route_zero_token = torch.div( + route_zero_position, + output_index.shape[1], + rounding_mode="floor", + ) + route_zero_weight = flat_weights.index_select(0, route_zero_position).float() + route_zero_grad = ( + grad_output.index_select(0, route_zero_token).float() * route_zero_weight.unsqueeze(1) + ).to(dtype=grad_routes.dtype) + grad_routes.narrow(0, 0, 1).copy_(route_zero_grad) + return + + valid_positions = torch.nonzero(output_index >= 0, as_tuple=False) + # If the returned input gradient aliases route_values, finish every + # probability gradient before clearing or overwriting that storage. + # Padded DeepEP batches contain -1 output indices, and aligned expert rows + # not referenced by a real token must receive a zero gradient. + if routes_alias_values and grad_weights is not None: + for start in range(0, valid_positions.shape[0], _BACKWARD_CHUNK_ROWS): + end = min(start + _BACKWARD_CHUNK_ROWS, valid_positions.shape[0]) + positions = valid_positions[start:end] + token_rows = positions[:, 0] + topk_columns = positions[:, 1] + route_rows = output_index[token_rows, topk_columns].to(dtype=torch.long) + token_grads = grad_output.index_select(0, token_rows) + selected_route_values = route_values.index_select(0, route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + grad_weights[token_rows, topk_columns] = weight_grads + + if grad_routes is not None and routes_alias_values: + grad_routes.zero_() + + for start in range(0, valid_positions.shape[0], _BACKWARD_CHUNK_ROWS): + end = min(start + _BACKWARD_CHUNK_ROWS, valid_positions.shape[0]) + positions = valid_positions[start:end] + token_rows = positions[:, 0] + topk_columns = positions[:, 1] + route_rows = output_index[token_rows, topk_columns].to(dtype=torch.long) + token_grads = grad_output.index_select(0, token_rows) + + if grad_routes is not None: + route_grads = (token_grads.float() * topk_weights[token_rows, topk_columns].float().unsqueeze(1)).to( + dtype=grad_routes.dtype + ) + grad_routes.index_copy_(0, route_rows, route_grads) + if grad_weights is not None and not routes_alias_values: + selected_route_values = route_values.index_select(0, route_rows) + weight_grads = (token_grads.float() * selected_route_values.float()).sum(dim=-1) + grad_weights[token_rows, topk_columns] = weight_grads + + +class _DeepGEMMMoEWithBF16Backward(torch.autograd.Function): + """FP8 grouped-DeepGEMM MoE forward with explicit BF16 backward.""" + + @staticmethod + def forward( + ctx, + permuted_local_hidden_states: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + module: torch.nn.Module, + layout: _MoELayout, + module_name: str, + *weights: torch.Tensor, + ) -> torch.Tensor: + if len(weights) != 2 * layout.num_local_experts: + raise RuntimeError( + f"{module_name} expected {2 * layout.num_local_experts} expert weights, got {len(weights)}" + ) + counts = _validate_routing_inputs( + permuted_local_hidden_states, + tokens_per_expert, + permuted_probs, + layout, + ) + output = _deepgemm_grouped_moe_forward( + module, + permuted_local_hidden_states, + tokens_per_expert, + permuted_probs, + layout=layout, + module_name=module_name, + validated_counts=counts, + ) + ctx.layout = layout + ctx.counts = counts + ctx.module_name = module_name + ctx.defer_router_probabilities = bool(getattr(module, "_vime_defer_router_probabilities", False)) + ctx.reuse_expert_input_for_grad = bool(getattr(module, "_vime_reuse_expert_input_for_grad", False)) + ctx.grad_workspace = getattr(module, _COMBINE_WORKSPACE_ATTR, None) + ctx.save_for_backward(permuted_local_hidden_states, permuted_probs, *weights) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + saved = ctx.saved_tensors + hidden_states = saved[0] + permuted_probs = saved[1] + weights = saved[2:] + layout: _MoELayout = ctx.layout + counts: tuple[int, ...] = ctx.counts + + fc1_weights = weights[: layout.num_local_experts] + fc2_weights = weights[layout.num_local_experts :] + if len(fc1_weights) != layout.num_local_experts or len(fc2_weights) != layout.num_local_experts: + raise RuntimeError(f"{ctx.module_name} saved expert weight count mismatch") + + needs = ctx.needs_input_grad + needs_hidden = needs[0] + needs_probs = needs[2] + needs_fc1_weights = needs[6 : 6 + layout.num_local_experts] + needs_fc2_weights = needs[6 + layout.num_local_experts :] + + grad_hidden = None + if needs_hidden: + workspace = ctx.grad_workspace + if workspace is None and ctx.reuse_expert_input_for_grad: + # DeepEP's expert-major input has no later reader. Each + # expert/chunk finishes recompute and wgrad before writing its + # dgrad, so that input storage can safely carry the gradient. + grad_hidden = hidden_states.detach() + elif workspace is None: + grad_hidden = torch.empty_like(hidden_states) + else: + required_bytes = hidden_states.numel() * hidden_states.element_size() + if required_bytes > workspace.numel(): + raise RuntimeError( + "Shared MoE backward workspace is too small: " + f"need {required_bytes} bytes for {tuple(hidden_states.shape)}, " + f"have {workspace.numel()} bytes; increase " + "VIME_DEEPGEMM_MOE_COMBINE_WORKSPACE_BYTES" + ) + if workspace.device != hidden_states.device: + raise RuntimeError( + "Shared MoE backward workspace is on " + f"{workspace.device}, input is on {hidden_states.device}" + ) + grad_hidden = workspace.narrow(0, 0, required_bytes).view(hidden_states.dtype).view_as(hidden_states) + grad_probs = torch.empty_like(permuted_probs) if needs_probs else None + grad_fc1_weights: list[torch.Tensor | None] = [None] * layout.num_local_experts + grad_fc2_weights: list[torch.Tensor | None] = [None] * layout.num_local_experts + grad_output = grad_output.contiguous().to(dtype=hidden_states.dtype) + probabilities = permuted_probs.reshape(-1, 1) + defer_router_probabilities = ctx.defer_router_probabilities + + if _use_grouped_bf16_backward( + hidden_states, + counts, + needs_fc1_weights, + needs_fc2_weights, + ): + grad_hidden, grad_probs, grad_fc1_weights, grad_fc2_weights = _grouped_expert_backward( + hidden_states=hidden_states, + permuted_probs=permuted_probs, + grad_output=grad_output, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + counts=counts, + layout=layout, + needs_hidden=needs_hidden, + needs_probs=needs_probs, + needs_fc1_weights=needs_fc1_weights, + needs_fc2_weights=needs_fc2_weights, + defer_router_probabilities=defer_router_probabilities, + grad_hidden=grad_hidden, + grad_probs=grad_probs, + ) + return ( + grad_hidden, + None, + grad_probs, + None, + None, + None, + *grad_fc1_weights, + *grad_fc2_weights, + ) + + offset = 0 + for expert_index, count in enumerate(counts): + fc1_weight = fc1_weights[expert_index] + fc2_weight = fc2_weights[expert_index] + needs_fc1_weight = needs_fc1_weights[expert_index] + needs_fc2_weight = needs_fc2_weights[expert_index] + if count == 0: + if needs_fc1_weight: + grad_fc1_weights[expert_index] = torch.zeros_like(fc1_weight) + if needs_fc2_weight: + grad_fc2_weights[expert_index] = torch.zeros_like(fc2_weight) + continue + + fc1_accumulator = torch.zeros_like(fc1_weight, dtype=torch.float32) if needs_fc1_weight else None + fc2_accumulator = torch.zeros_like(fc2_weight, dtype=torch.float32) if needs_fc2_weight else None + + for chunk_start in range(0, count, _BACKWARD_CHUNK_ROWS): + chunk_end = min(chunk_start + _BACKWARD_CHUNK_ROWS, count) + global_start = offset + chunk_start + global_end = offset + chunk_end + hidden = hidden_states[global_start:global_end] + grad = grad_output[global_start:global_end] + probability = probabilities[global_start:global_end] + + gate_up = _deepgemm_bf16_gemm_nt(hidden, fc1_weight) + gate, up = gate_up.chunk(2, dim=-1) + gate_f = gate.float() + up_f = up.float() + silu_gate = F.silu(gate_f) + down_input = (silu_gate * up_f).to(dtype=hidden_states.dtype) + + if needs_probs and not defer_router_probabilities: + down_output = _deepgemm_bf16_gemm_nt(down_input, fc2_weight) + grad_probs_chunk = _router_probability_grad_fp32_chunked( + grad, + down_output, + ) + grad_probs[global_start:global_end].copy_( + grad_probs_chunk.reshape_as(permuted_probs[global_start:global_end]).to( + dtype=permuted_probs.dtype + ) + ) + + if defer_router_probabilities: + grad_down_output = grad + else: + grad_down_output = (grad.float() * probability.float()).to(dtype=hidden_states.dtype) + grad_down_input = _deepgemm_bf16_gemm_nn( + grad_down_output, + fc2_weight, + ) + if fc2_accumulator is not None: + fc2_accumulator.add_( + _deepgemm_bf16_gemm_tn( + grad_down_output, + down_input, + ) + ) + + grad_down_input_f = grad_down_input.float() + sigmoid_gate = torch.sigmoid(gate_f) + grad_gate = grad_down_input_f * up_f * sigmoid_gate * (1.0 + gate_f * (1.0 - sigmoid_gate)) + grad_up = grad_down_input_f * silu_gate + grad_gate_up = torch.cat([grad_gate, grad_up], dim=-1).to(dtype=hidden_states.dtype) + + if fc1_accumulator is not None: + fc1_accumulator.add_(_deepgemm_bf16_gemm_tn(grad_gate_up, hidden)) + # Keep this after the final read from ``hidden`` because the + # DeepEP path may reuse that storage for grad_hidden. + if needs_hidden: + grad_hidden[global_start:global_end].copy_(_deepgemm_bf16_gemm_nn(grad_gate_up, fc1_weight)) + + if fc1_accumulator is not None: + grad_fc1_weights[expert_index] = _sum_to_parameter_dtype( + fc1_accumulator, + fc1_weight, + ) + if fc2_accumulator is not None: + grad_fc2_weights[expert_index] = _sum_to_parameter_dtype( + fc2_accumulator, + fc2_weight, + ) + offset += count + + return ( + grad_hidden, + None, + None if defer_router_probabilities else grad_probs, + None, + None, + None, + *grad_fc1_weights, + *grad_fc2_weights, + ) + + +def _configure_batch_invariant(deep_gemm: Any) -> bool: + enabled = os.environ.get("VLLM_BATCH_INVARIANT", "").lower() in { + "1", + "true", + "yes", + "on", + } + setter = getattr(deep_gemm, "set_batch_invariant", None) + if setter is None: + raise RuntimeError("deep_gemm.set_batch_invariant is unavailable") + setter(enabled) + getter = getattr(deep_gemm, "get_batch_invariant", None) + if enabled and (getter is None or not getter()): + raise RuntimeError( + "VLLM_BATCH_INVARIANT=1, but the Megatron actor's " + "DeepGEMM runtime did not enable batch-invariant kernels" + ) + return enabled + + +def _load_deepgemm_ops() -> _DeepGEMMOps: + """Load CUDA-only dependencies lazily so CPU tests can mock this boundary.""" + import deep_gemm + from vllm.model_executor.layers.quantization.utils.fp8_utils import per_token_group_quant_fp8 + from vllm.utils import deep_gemm as vllm_deep_gemm + + from vime.backends.megatron_utils.kernels.fp8_kernel import blockwise_cast_to_fp8_triton + + _configure_batch_invariant(deep_gemm) + scale_ue8m0 = vllm_deep_gemm.is_deep_gemm_e8m0_used() + m_alignment = int(vllm_deep_gemm.get_mk_alignment_for_contiguous_layout()[0]) + if m_alignment != _GROUPED_M_ALIGNMENT: + raise RuntimeError( + f"Unexpected DeepGEMM contiguous grouped M alignment: {m_alignment} != {_GROUPED_M_ALIGNMENT}" + ) + + if not scale_ue8m0: + # Hopper (sm90): FP32 block scales; weights cast with the Triton kernel + # and activation/weight scales TMA-aligned as a separate step. Unchanged. + return _DeepGEMMOps( + quantize_weight=blockwise_cast_to_fp8_triton, + quantize_activation=per_token_group_quant_fp8, + align_input_scale=vllm_deep_gemm.get_col_major_tma_aligned_tensor, + grouped_gemm=vllm_deep_gemm.m_grouped_fp8_gemm_nt_contiguous, + silu_and_mul=_vllm_silu_and_mul, + scale_ue8m0=False, + need_tma_aligned_scales=True, + transform_weight_scale=None, + ) + + # Blackwell (sm100+): VLLM uses UE8M0 power-of-two block scales. Activation + # scales are produced column-major and TMA-aligned directly by the native + # quantization helper. + def _quantize_weight_ue8m0(weight: torch.Tensor, block: tuple[int, int]): + # Mirror the Hopper op signature (weight, (block_n, block_k)); UE8M0 quant + # requires a [128, 128] block and a BF16 input. + return vllm_deep_gemm.per_block_cast_to_fp8( + weight, + block_size=(int(block[0]), int(block[1])), + ) + + return _DeepGEMMOps( + quantize_weight=_quantize_weight_ue8m0, + quantize_activation=per_token_group_quant_fp8, + align_input_scale=vllm_deep_gemm.get_col_major_tma_aligned_tensor, + grouped_gemm=vllm_deep_gemm.m_grouped_fp8_gemm_nt_contiguous, + silu_and_mul=_vllm_silu_and_mul, + scale_ue8m0=True, + need_tma_aligned_scales=False, + transform_weight_scale=None, + ) + + +def _get_expert_weights( + grouped_linear: torch.nn.Module, + *, + num_local_experts: int, + expected_shape: tuple[int, int], + module_name: str, +) -> list[torch.Tensor]: + weights = [] + for expert_index in range(num_local_experts): + attr_name = f"weight{expert_index}" + weight = getattr(grouped_linear, attr_name, None) + if not isinstance(weight, torch.Tensor): + raise RuntimeError(f"{module_name} is missing tensor parameter {attr_name}") + if tuple(weight.shape) != expected_shape: + raise RuntimeError( + f"{module_name}.{attr_name} has shape {tuple(weight.shape)}, " + f"expected {expected_shape} ([out_features, in_features])" + ) + weights.append(weight) + return weights + + +def _validate_parallelism() -> None: + tp_size = parallel_state.get_tensor_model_parallel_world_size() + if tp_size != 1: + raise RuntimeError(f"TEGroupedMLP DeepGEMM alignment requires tensor model parallel size 1, got {tp_size}") + expert_tp_size = parallel_state.get_expert_tensor_parallel_world_size() + if expert_tp_size != 1: + raise RuntimeError( + f"TEGroupedMLP DeepGEMM alignment requires expert tensor parallel size 1, got {expert_tp_size}" + ) + + +def _validate_te_grouped_mlp(module: torch.nn.Module, module_name: str) -> _MoELayout: + if type(module).__name__ != "TEGroupedMLP": + raise RuntimeError( + f"DeepGEMM MoE target {module_name} has unsupported class {type(module).__name__}; expected TEGroupedMLP" + ) + + config = getattr(module, "config", None) + if config is None: + raise RuntimeError(f"DeepGEMM MoE target {module_name} has no TransformerConfig") + if getattr(config, "add_bias_linear", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} must be bias-free") + if not getattr(config, "gated_linear_unit", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} must use a gated linear unit") + if getattr(module, "activation_func", None) is not F.silu: + raise RuntimeError(f"DeepGEMM MoE target {module_name} must use SiLU") + if getattr(config, "moe_apply_probs_on_input", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} cannot use moe_apply_probs_on_input") + if getattr(config, "fp8", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} cannot also enable Megatron FP8") + if getattr(config, "qat", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} cannot enable QAT") + if getattr(config, "swiglu_clamp_limit", None) is not None: + raise RuntimeError( + f"DeepGEMM MoE target {module_name} has a SwiGLU clamp, " + "which the VLLM DeepGEMM activation used here does not reproduce" + ) + + num_local_experts = int(getattr(module, "num_local_experts", 0)) + hidden_size = int(getattr(config, "hidden_size", 0)) + ffn_hidden_size = int(getattr(config, "moe_ffn_hidden_size", 0)) + if min(num_local_experts, hidden_size, ffn_hidden_size) <= 0: + raise RuntimeError( + f"DeepGEMM MoE target {module_name} has invalid layout: " + f"experts={num_local_experts}, hidden={hidden_size}, ffn={ffn_hidden_size}" + ) + if hidden_size % _BLOCK_SIZE or ffn_hidden_size % _BLOCK_SIZE: + raise RuntimeError( + f"DeepGEMM MoE target {module_name} requires hidden and MoE FFN sizes " + f"divisible by {_BLOCK_SIZE}, got {hidden_size} and {ffn_hidden_size}" + ) + + fc1 = getattr(module, "linear_fc1", None) + fc2 = getattr(module, "linear_fc2", None) + if type(fc1).__name__ != "TEColumnParallelGroupedLinear": + raise RuntimeError(f"{module_name}.linear_fc1 must be TEColumnParallelGroupedLinear, got {type(fc1).__name__}") + if type(fc2).__name__ != "TERowParallelGroupedLinear": + raise RuntimeError(f"{module_name}.linear_fc2 must be TERowParallelGroupedLinear, got {type(fc2).__name__}") + if int(getattr(fc1, "num_gemms", -1)) != num_local_experts: + raise RuntimeError(f"{module_name}.linear_fc1 num_gemms does not match local experts") + if int(getattr(fc2, "num_gemms", -1)) != num_local_experts: + raise RuntimeError(f"{module_name}.linear_fc2 num_gemms does not match local experts") + if getattr(fc1, "use_bias", False) or getattr(fc2, "use_bias", False): + raise RuntimeError(f"DeepGEMM MoE target {module_name} grouped linears must be bias-free") + if getattr(fc1, "_vime_deepgemm_forward_wrapped", False) or getattr(fc2, "_vime_deepgemm_forward_wrapped", False): + raise RuntimeError( + f"DeepGEMM MoE target {module_name} has an individually wrapped expert linear; " + "remove that wrapper before installing the whole-MLP hook" + ) + + layout = _MoELayout( + num_local_experts=num_local_experts, + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + ) + fc1_weights = _get_expert_weights( + fc1, + num_local_experts=num_local_experts, + expected_shape=layout.fc1_weight_shape, + module_name=f"{module_name}.linear_fc1", + ) + fc2_weights = _get_expert_weights( + fc2, + num_local_experts=num_local_experts, + expected_shape=layout.fc2_weight_shape, + module_name=f"{module_name}.linear_fc2", + ) + named_weights = [(f"linear_fc1.weight{i}", weight) for i, weight in enumerate(fc1_weights)] + named_weights.extend((f"linear_fc2.weight{i}", weight) for i, weight in enumerate(fc2_weights)) + for weight_name, weight in named_weights: + if weight.dtype != torch.bfloat16: + raise RuntimeError(f"{module_name}.{weight_name} must be BF16, got {weight.dtype}") + return layout + + +def _validate_routing_inputs( + hidden_states: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + layout: _MoELayout, +) -> tuple[int, ...]: + if hidden_states.dtype != torch.bfloat16: + raise RuntimeError(f"DeepGEMM MoE alignment requires BF16 hidden states, got {hidden_states.dtype}") + if hidden_states.ndim != 2 or hidden_states.shape[1] != layout.hidden_size: + raise RuntimeError( + "DeepGEMM MoE hidden-state shape mismatch: " + f"got {tuple(hidden_states.shape)}, expected [tokens, {layout.hidden_size}]" + ) + if not isinstance(tokens_per_expert, torch.Tensor): + raise RuntimeError("tokens_per_expert must be a tensor") + if tokens_per_expert.ndim != 1 or tokens_per_expert.numel() != layout.num_local_experts: + raise RuntimeError( + "tokens_per_expert must have one entry per local expert: " + f"got {tuple(tokens_per_expert.shape)}, expected [{layout.num_local_experts}]" + ) + if tokens_per_expert.dtype not in { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + }: + raise RuntimeError(f"tokens_per_expert must have an integer dtype, got {tokens_per_expert.dtype}") + + counts = tuple(int(value) for value in tokens_per_expert.detach().cpu().tolist()) + if any(value < 0 for value in counts): + raise RuntimeError(f"tokens_per_expert contains a negative count: {counts}") + if sum(counts) != hidden_states.shape[0]: + raise RuntimeError( + "tokens_per_expert sum does not match permuted hidden-state rows: " + f"{sum(counts)} != {hidden_states.shape[0]}" + ) + if not isinstance(permuted_probs, torch.Tensor) or not permuted_probs.is_floating_point(): + raise RuntimeError("permuted_probs must be a floating-point tensor") + if permuted_probs.ndim not in (1, 2) or permuted_probs.numel() != hidden_states.shape[0]: + raise RuntimeError( + "permuted_probs must contain one scalar per permuted row: " + f"got {tuple(permuted_probs.shape)}, expected [{hidden_states.shape[0]}]" + ) + if permuted_probs.ndim == 2 and permuted_probs.shape[1] != 1: + raise RuntimeError(f"2D permuted_probs must have shape [tokens, 1], got {tuple(permuted_probs.shape)}") + if permuted_probs.device != hidden_states.device: + raise RuntimeError( + "permuted_probs and hidden states must be on the same device: " + f"{permuted_probs.device} != {hidden_states.device}" + ) + return counts + + +def _quantize_grouped_weights( + grouped_linear: torch.nn.Module, + *, + expected_shape: tuple[int, int], + layout: _MoELayout, + input_device: torch.device, + module_name: str, + ops: _DeepGEMMOps, + expert_start: int = 0, + expert_end: int | None = None, + grouped_qweight_workspace: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + weights = _get_expert_weights( + grouped_linear, + num_local_experts=layout.num_local_experts, + expected_shape=expected_shape, + module_name=module_name, + ) + expert_end = layout.num_local_experts if expert_end is None else expert_end + if not 0 <= expert_start < expert_end <= layout.num_local_experts: + raise RuntimeError( + f"Invalid {module_name} expert range [{expert_start}, {expert_end}) " + f"for {layout.num_local_experts} local experts" + ) + weights = weights[expert_start:expert_end] + num_group_experts = expert_end - expert_start + + grouped_qweight = None + grouped_scale = None + expected_scale_shape = ( + (expected_shape[0] + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + (expected_shape[1] + _BLOCK_SIZE - 1) // _BLOCK_SIZE, + ) + for local_expert_index, weight in enumerate(weights): + expert_index = expert_start + local_expert_index + if weight.dtype != torch.bfloat16: + raise RuntimeError(f"{module_name}.weight{expert_index} must be BF16, got {weight.dtype}") + if weight.device != input_device: + raise RuntimeError( + f"{module_name}.weight{expert_index} and hidden states must be on the same " + f"device: {weight.device} != {input_device}" + ) + qweight, scale = ops.quantize_weight( + weight.detach().contiguous(), + (_BLOCK_SIZE, _BLOCK_SIZE), + ) + if tuple(qweight.shape) != expected_shape: + raise RuntimeError( + f"quantized {module_name}.weight{expert_index} has shape " + f"{tuple(qweight.shape)}, expected {expected_shape}" + ) + if tuple(scale.shape) != expected_scale_shape or (not ops.scale_ue8m0 and scale.dtype != torch.float32): + raise RuntimeError( + f"quantized {module_name}.weight{expert_index} scale has shape/dtype " + f"{tuple(scale.shape)}/{scale.dtype}, expected " + f"{expected_scale_shape}/{torch.float32}" + ) + if qweight.device != input_device or scale.device != input_device: + raise RuntimeError(f"quantized {module_name}.weight{expert_index} is on the wrong device") + + if grouped_qweight is None: + grouped_shape = (num_group_experts, *expected_shape) + grouped_qweight = _storage_prefix_view( + grouped_qweight_workspace, + grouped_shape, + qweight.dtype, + ) + if grouped_qweight is None: + grouped_qweight = qweight.new_empty(grouped_shape) + grouped_scale = scale.new_empty((num_group_experts, *expected_scale_shape)) + elif qweight.dtype != grouped_qweight.dtype: + raise RuntimeError(f"quantized {module_name} expert weights have inconsistent dtypes") + grouped_qweight[local_expert_index].copy_(qweight) + grouped_scale[local_expert_index].copy_(scale) + # Do not retain the previous per-expert quantization while the next + # expert is being quantized. At GLM-750B dimensions each FC1 value is + # 24 MiB, so even one stale loop value matters at this peak. + del qweight, scale + + assert grouped_qweight is not None and grouped_scale is not None + if ops.scale_ue8m0: + # Blackwell: the VLLM rollout stores UE8M0 MoE weights by quantizing the + # BF16 experts to FP8 and then *requantizing* the grouped weight in + # process_weights_after_loading (requant_block_scale_ue8m0_for_deepgemm -> + # requant_weight_ue8m0). That requant is lossy, so a single quant here + # would not bit-match the rollout. Replicate quant -> requant on the + # grouped [E, N, K] weight to align to ~e-7. + from vllm.model_executor.layers.quantization.utils.fp8_utils import requant_weight_ue8m0_inplace + + requant_weight_ue8m0_inplace(grouped_qweight, grouped_scale, (_BLOCK_SIZE, _BLOCK_SIZE)) + return grouped_qweight, grouped_scale + + +def _storage_prefix_view( + workspace: torch.Tensor | None, + shape: tuple[int, ...], + dtype: torch.dtype, +) -> torch.Tensor | None: + """Return a typed prefix of dead contiguous tensor storage, if it fits.""" + if workspace is None or not workspace.is_contiguous(): + return None + required_elements = math.prod(shape) + required_bytes = required_elements * torch.empty((), dtype=dtype).element_size() + available_bytes = workspace.numel() * workspace.element_size() + if required_bytes > available_bytes: + return None + raw_workspace = workspace.view(torch.uint8).reshape(-1) + return raw_workspace.narrow(0, 0, required_bytes).view(dtype).view(shape) + + +def _quantize_activation( + value: torch.Tensor, + ops: _DeepGEMMOps, +) -> tuple[torch.Tensor, torch.Tensor]: + ue8m0 = ops.scale_ue8m0 + qvalue, scale = ops.quantize_activation( + value.detach().contiguous(), + _BLOCK_SIZE, + column_major_scales=ue8m0, + tma_aligned_scales=ue8m0, + use_ue8m0=ue8m0, + ) + if qvalue.shape != value.shape: + raise RuntimeError(f"quantized activation shape mismatch: {tuple(qvalue.shape)} != {tuple(value.shape)}") + if not ue8m0: + # Hopper: FP32 row-major block scales, TMA-aligned as a separate step. + expected_scale_shape = (value.shape[0], value.shape[1] // _BLOCK_SIZE) + if tuple(scale.shape) != expected_scale_shape or scale.dtype != torch.float32: + raise RuntimeError( + "quantized activation scale shape/dtype mismatch: " + f"{tuple(scale.shape)}/{scale.dtype} != {expected_scale_shape}/{torch.float32}" + ) + return qvalue, ops.align_input_scale(scale) + # Blackwell: the native quantization helper already produced column-major, + # TMA-aligned UE8M0 scales, so no separate alignment step is needed. + return qvalue, scale + + +def _build_m_indices( + counts: tuple[int, ...], + *, + device: torch.device, +) -> torch.Tensor: + # Counts are already available as Python integers for the per-expert + # grouped launches. Filling the device output directly avoids both a + # blocking host-to-device upload of a newly constructed repeats tensor and + # repeat_interleave's data-dependent shape path. + output = torch.empty(sum(counts), dtype=torch.int32, device=device) + start = 0 + for expert, count in enumerate(counts): + if count: + output.narrow(0, start, count).fill_(expert) + start += count + return output + + +def _experts_per_forward_group(num_local_experts: int) -> int: + configured = os.environ.get("VIME_DEEPGEMM_MOE_EXPERTS_PER_GROUP") + if configured is None: + batch_invariant = os.environ.get("VLLM_BATCH_INVARIANT", "").lower() in { + "1", + "true", + "yes", + "on", + } + return min(_DEFAULT_EXPERTS_PER_GROUP, num_local_experts) if batch_invariant else num_local_experts + try: + experts_per_group = int(configured) + except ValueError as exc: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_EXPERTS_PER_GROUP must be a positive integer, " f"got {configured!r}" + ) from exc + if experts_per_group <= 0: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_EXPERTS_PER_GROUP must be a positive integer, " f"got {experts_per_group}" + ) + return min(experts_per_group, num_local_experts) + + +def _sort_chunks_into( + input_: torch.Tensor, + split_sizes: torch.Tensor, + sorted_idxs: torch.Tensor, + output: torch.Tensor, +) -> torch.Tensor: + """Megatron ``sort_chunks_by_idxs`` ordering with caller-owned storage. + + Keep the chunk metadata on the input device. Calling ``.tolist()`` on the + CUDA split/order tensors serializes the stream once per MoE permutation. + The destination-row map below expresses the same chunk-stable permutation + with device operations and writes directly into the caller's workspace. + """ + if output.shape != input_.shape or output.dtype != input_.dtype or output.device != input_.device: + raise RuntimeError( + "Preallocated MoE combine buffer must match its input: " + f"input={tuple(input_.shape)}/{input_.dtype}/{input_.device}, " + f"output={tuple(output.shape)}/{output.dtype}/{output.device}" + ) + if split_sizes.ndim != 1 or sorted_idxs.ndim != 1: + raise RuntimeError( + "MoE chunk sizes/order must be one-dimensional: " + f"{tuple(split_sizes.shape)} and {tuple(sorted_idxs.shape)}" + ) + if split_sizes.numel() != sorted_idxs.numel(): + raise RuntimeError("MoE chunk sizes/order length mismatch: " f"{split_sizes.numel()} != {sorted_idxs.numel()}") + if input_.shape[0] == 0: + return output + + # Without fused permutation Megatron has already copied this metadata to + # CPU and synchronized its side stream. Preserve the cheap host path in + # that case; moving the metadata back to CUDA would add a new transfer. + # The CUDA branch below is what removes the hot-path synchronization when + # moe_permute_fusion keeps both tensors on device. + if not split_sizes.is_cuda and not sorted_idxs.is_cuda: + chunks = torch.split(input_, split_sizes.tolist(), dim=0) + output_offset = 0 + for index in sorted_idxs.tolist(): + chunk = chunks[index] + chunk_rows = chunk.shape[0] + output.narrow(0, output_offset, chunk_rows).copy_(chunk) + output_offset += chunk_rows + if output_offset != input_.shape[0]: + raise RuntimeError(f"Preallocated MoE combine copied {output_offset} rows, " f"expected {input_.shape[0]}") + return output + + device = input_.device + sizes = split_sizes.to(device=device, dtype=torch.long) + order = sorted_idxs.to(device=device, dtype=torch.long) + num_chunks = sizes.numel() + + # The metadata is generated by Megatron's dispatcher. Keep defensive + # checks asynchronous for CUDA callers so they do not recreate the host + # synchronization this path is meant to remove. + torch._assert_async(torch.all(sizes >= 0), "MoE chunk sizes cannot be negative") + torch._assert_async( + sizes.sum() == input_.shape[0], + "MoE chunk sizes must sum to the input row count", + ) + torch._assert_async( + torch.all((order >= 0) & (order < num_chunks)), + "MoE chunk order contains an out-of-range index", + ) + + source_chunks = torch.repeat_interleave( + torch.arange(num_chunks, device=device, dtype=torch.long), + sizes, + output_size=input_.shape[0], + ) + source_starts = torch.cumsum(sizes, dim=0) - sizes + within_chunk = torch.arange(input_.shape[0], device=device, dtype=torch.long) - source_starts.index_select( + 0, source_chunks + ) + + sorted_sizes = sizes.index_select(0, order) + sorted_starts = torch.cumsum(sorted_sizes, dim=0) - sorted_sizes + destination_starts = torch.empty_like(sorted_starts) + destination_starts.scatter_(0, order, sorted_starts) + destination_rows = destination_starts.index_select(0, source_chunks) + within_chunk + output.index_copy_(0, destination_rows, input_) + return output + + +def _wrap_preallocated_combine_preprocess(dispatcher: torch.nn.Module) -> bool: + """Consume the expert's early-allocated combine buffer without a peak allocation.""" + if getattr(dispatcher, "_vime_preallocated_combine_wrapped", False): + return False + if int(getattr(dispatcher, "tp_size", 1)) != 1: + raise RuntimeError("Preallocated MoE combine currently requires tensor parallel size 1") + if bool(getattr(dispatcher, "drop_and_pad", False)): + raise RuntimeError("Preallocated MoE combine does not support moe_expert_capacity_factor") + + original_combine_preprocess = dispatcher.combine_preprocess + + def combine_preprocess( + patched_dispatcher: torch.nn.Module, + hidden_states: torch.Tensor, + ): + combine_buffer = getattr(hidden_states, _PREALLOCATED_COMBINE_BUFFER_ATTR, None) + if combine_buffer is None: + return original_combine_preprocess(hidden_states) + if int(patched_dispatcher.num_local_experts) <= 1: + return hidden_states + combined = _sort_chunks_into( + hidden_states, + patched_dispatcher.num_global_tokens_per_local_expert.T.ravel(), + patched_dispatcher.restore_output_by_local_experts, + combine_buffer, + ) + setattr(combined, _PREALLOCATED_TOKEN_COMBINE_ATTR, True) + return combined + + dispatcher.combine_preprocess = types.MethodType(combine_preprocess, dispatcher) + dispatcher._vime_preallocated_combine_wrapped = True + return True + + +def _wrap_preallocated_dispatch_postprocess(dispatcher: torch.nn.Module) -> bool: + """Use the shared workspace for no-grad dispatch permutation. + + The all-to-all output is kept as the later combine destination. With TP=1, + shared-expert overlap does not consume this tensor's values after + ``linear_fc1_forward_and_act``; it only tags the tensor for backward order. + """ + if getattr(dispatcher, "_vime_preallocated_dispatch_wrapped", False): + return False + if int(getattr(dispatcher, "tp_size", 1)) != 1: + raise RuntimeError("Preallocated MoE dispatch currently requires tensor parallel size 1") + if bool(getattr(dispatcher, "drop_and_pad", False)): + raise RuntimeError("Preallocated MoE dispatch does not support moe_expert_capacity_factor") + if bool(getattr(getattr(dispatcher, "config", None), "moe_permute_fusion", False)): + raise RuntimeError("Preallocated MoE dispatch currently requires moe_permute_fusion disabled") + + original_dispatch_postprocess = dispatcher.dispatch_postprocess + + def dispatch_postprocess( + patched_dispatcher: torch.nn.Module, + global_input_tokens: torch.Tensor, + global_probs: torch.Tensor, + ): + if int(patched_dispatcher.num_local_experts) <= 1: + return original_dispatch_postprocess(global_input_tokens, global_probs) + if torch.is_grad_enabled(): + dispatched_input, tokens_per_expert, permuted_probs = original_dispatch_postprocess( + global_input_tokens, global_probs + ) + # The dispatch all-to-all output is no longer read after permutation: + # AllToAllBackward saves only the group/splits, CatBackward saves no + # input values, and shared experts consume cached_fc1_input instead. + # Keep its storage as the destination for the inverse expert sort. + # This removes the otherwise full-sized torch.cat allocation during + # checkpoint recomputation without overwriting the expert input that + # the explicit BF16 backward must retain. + setattr( + dispatched_input, + _PREALLOCATED_COMBINE_BUFFER_ATTR, + global_input_tokens, + ) + return dispatched_input, tokens_per_expert, permuted_probs + workspace_output = _combine_workspace_view( + patched_dispatcher, + global_input_tokens, + ) + if workspace_output is None: + return original_dispatch_postprocess(global_input_tokens, global_probs) + + if patched_dispatcher.shared_experts is not None: + patched_dispatcher.shared_experts.linear_fc1_forward_and_act(global_input_tokens) + patched_dispatcher.tokens_per_expert = patched_dispatcher._maybe_dtoh_and_synchronize( + "before_permutation_2", + patched_dispatcher.tokens_per_expert, + ) + split_sizes = patched_dispatcher.num_global_tokens_per_local_expert.ravel() + workspace_output = _sort_chunks_into( + global_input_tokens, + split_sizes, + patched_dispatcher.sort_input_by_local_experts, + workspace_output, + ) + sorted_probs = _sort_chunks_into( + global_probs, + split_sizes, + patched_dispatcher.sort_input_by_local_experts, + torch.empty_like(global_probs), + ) + setattr( + workspace_output, + _PREALLOCATED_COMBINE_BUFFER_ATTR, + global_input_tokens, + ) + tokens_per_expert = patched_dispatcher._maybe_dtoh_and_synchronize( + "before_finish", + patched_dispatcher.tokens_per_expert, + ) + patched_dispatcher.tokens_per_expert = None + return workspace_output, tokens_per_expert, sorted_probs + + dispatcher.dispatch_postprocess = types.MethodType(dispatch_postprocess, dispatcher) + dispatcher._vime_preallocated_dispatch_wrapped = True + return True + + +def _wrap_preallocated_token_combine(dispatcher: torch.nn.Module) -> bool: + """Write no-grad expert all-to-all results back into the shared workspace.""" + if getattr(dispatcher, "_vime_preallocated_token_combine_wrapped", False): + return False + original_token_combine = dispatcher.token_combine + + def token_combine( + patched_dispatcher: torch.nn.Module, + hidden_states: torch.Tensor, + *args, + **kwargs, + ): + if torch.is_grad_enabled() or not bool(getattr(hidden_states, _PREALLOCATED_TOKEN_COMBINE_ATTR, False)): + return original_token_combine(hidden_states, *args, **kwargs) + if patched_dispatcher.ep_group.size() == 1: + return hidden_states + output_rows = sum(patched_dispatcher.input_splits) + output = _combine_workspace_view( + patched_dispatcher, + hidden_states, + shape=(output_rows, *hidden_states.shape[1:]), + ) + if output is None: + return original_token_combine(hidden_states, *args, **kwargs) + torch.distributed.all_to_all_single( + output, + hidden_states, + output_split_sizes=patched_dispatcher.input_splits, + input_split_sizes=patched_dispatcher.output_splits, + group=patched_dispatcher.ep_group, + ) + return output + + dispatcher.token_combine = types.MethodType(token_combine, dispatcher) + dispatcher._vime_preallocated_token_combine_wrapped = True + return True + + +def _combine_workspace_bytes() -> int | None: + configured = os.environ.get("VIME_DEEPGEMM_MOE_COMBINE_WORKSPACE_BYTES") + if configured is None: + return None + try: + workspace_bytes = int(configured) + except ValueError as exc: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_COMBINE_WORKSPACE_BYTES must be a positive integer, " f"got {configured!r}" + ) from exc + if workspace_bytes <= 0: + raise RuntimeError( + "VIME_DEEPGEMM_MOE_COMBINE_WORKSPACE_BYTES must be a positive integer, " f"got {workspace_bytes}" + ) + return workspace_bytes + + +def _combine_workspace_view( + module: torch.nn.Module, + hidden_states: torch.Tensor, + *, + shape: tuple[int, ...] | None = None, +) -> torch.Tensor | None: + workspace = getattr(module, _COMBINE_WORKSPACE_ATTR, None) + if workspace is None: + return None + output_shape = tuple(hidden_states.shape) if shape is None else shape + required_elements = math.prod(output_shape) + required_bytes = required_elements * hidden_states.element_size() + if required_bytes > workspace.numel(): + raise RuntimeError( + "Shared MoE combine workspace is too small: " + f"need {required_bytes} bytes for {output_shape}, " + f"have {workspace.numel()} bytes; increase " + "VIME_DEEPGEMM_MOE_COMBINE_WORKSPACE_BYTES" + ) + if workspace.device != hidden_states.device: + raise RuntimeError( + f"Shared MoE combine workspace is on {workspace.device}, input is on {hidden_states.device}" + ) + return workspace.narrow(0, 0, required_bytes).view(hidden_states.dtype).view(output_shape) + + +def _moe_recompute_scratch_view( + module: torch.nn.Module, + reference: torch.Tensor, + shape: tuple[int, ...], +) -> torch.Tensor | None: + """Return a temporary tensor backed by the shared MoE workspace when it fits. + + This is used only by the gradient-enabled checkpoint recomputation. The + initial no-grad forward keeps its dispatched expert input in the same + workspace, so it must not use this scratch view. During recomputation the + workspace otherwise remains idle until the custom MoE backward writes + ``grad_hidden`` into it. + """ + workspace = getattr(module, _COMBINE_WORKSPACE_ATTR, None) + if workspace is None or workspace.device != reference.device: + return None + required_bytes = math.prod(shape) * reference.element_size() + if required_bytes > workspace.numel(): + return None + return workspace.narrow(0, 0, required_bytes).view(reference.dtype).view(shape) + + +def _pad_expert_rows( + value: torch.Tensor, + counts: tuple[int, ...], + *, + output: torch.Tensor | None = None, +) -> tuple[torch.Tensor, tuple[int, ...], torch.Tensor | None]: + """Pad every expert segment to the DeepGEMM contiguous-layout M alignment.""" + padded_counts = tuple( + ((count + _GROUPED_M_ALIGNMENT - 1) // _GROUPED_M_ALIGNMENT) * _GROUPED_M_ALIGNMENT if count else 0 + for count in counts + ) + if padded_counts == counts: + return value, padded_counts, None + + padded_shape = (sum(padded_counts), value.shape[1]) + if output is None: + padded_value = value.new_zeros(padded_shape) + else: + if tuple(output.shape) != padded_shape: + raise RuntimeError( + "Preallocated expert-row padding shape mismatch: " f"{tuple(output.shape)} != {padded_shape}" + ) + if output.dtype != value.dtype or output.device != value.device: + raise RuntimeError( + "Preallocated expert-row padding dtype/device mismatch: " + f"{output.dtype}/{output.device} != {value.dtype}/{value.device}" + ) + padded_value = output + padded_value.zero_() + valid_ranges = [] + padded_start = 0 + for count, padded_count in zip(counts, padded_counts, strict=True): + if count: + valid_ranges.append( + torch.arange( + padded_start, + padded_start + count, + device=value.device, + dtype=torch.long, + ) + ) + padded_start += padded_count + valid_rows = torch.cat(valid_ranges) + padded_value.index_copy_(0, valid_rows, value) + return padded_value, padded_counts, valid_rows + + +def _deepgemm_grouped_moe_forward( + module: torch.nn.Module, + hidden_states: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + *, + layout: _MoELayout, + module_name: str, + reuse_input_buffer: bool = False, + validated_counts: tuple[int, ...] | None = None, +) -> torch.Tensor: + """Compute the VLLM-style contiguous grouped-MoE forward without autograd.""" + counts = validated_counts + if counts is None: + counts = _validate_routing_inputs( + hidden_states, + tokens_per_expert, + permuted_probs, + layout, + ) + num_tokens = hidden_states.shape[0] + if num_tokens == 0: + return hidden_states.new_empty((0, layout.hidden_size)) + + ops = _load_deepgemm_ops() + experts_per_group = _experts_per_forward_group(layout.num_local_experts) + # Checkpointed forward runs under no_grad and does not need the dispatched + # expert-major input after each expert has consumed it. Reusing that + # storage removes one full [routed_tokens, hidden] allocation before + # Megatron's equally large combine permutation. Gradient-enabled + # recomputation keeps the ordinary separate output because the custom + # backward must save the original expert inputs. + shared_workspace = getattr(module, _COMBINE_WORKSPACE_ATTR, None) + direct_group_outputs = not reuse_input_buffer and shared_workspace is None + output_storage = None + if direct_group_outputs: + # During gradient-enabled checkpoint recomputation the expert input + # must remain live for the custom backward. Previously this allocated + # both the final [routes, hidden] output and another group-sized output + # (1.6 GiB each for the EP32/8K workload). Reserve only enough padding + # after the compact prefix for the largest current group, write the + # grouped GEMM there, then compact it in place. The returned prefix + # owns the slightly larger storage, but no second routed output exists. + compact_offset = 0 + output_capacity = num_tokens + for capacity_start in range(0, layout.num_local_experts, experts_per_group): + capacity_end = min( + capacity_start + experts_per_group, + layout.num_local_experts, + ) + capacity_counts = counts[capacity_start:capacity_end] + padded_capacity = sum( + ((count + _GROUPED_M_ALIGNMENT - 1) // _GROUPED_M_ALIGNMENT * _GROUPED_M_ALIGNMENT if count else 0) + for count in capacity_counts + ) + output_capacity = max( + output_capacity, + compact_offset + padded_capacity, + ) + compact_offset += sum(capacity_counts) + output_storage = hidden_states.new_empty((output_capacity, layout.hidden_size)) + output = output_storage.narrow(0, 0, num_tokens) + else: + output = hidden_states if reuse_input_buffer else hidden_states.new_empty((num_tokens, layout.hidden_size)) + defer_router_probabilities = bool(getattr(module, "_vime_defer_router_probabilities", False)) + + token_offset = 0 + for expert_start in range(0, layout.num_local_experts, experts_per_group): + expert_end = min(expert_start + experts_per_group, layout.num_local_experts) + group_counts = counts[expert_start:expert_end] + group_tokens = sum(group_counts) + group_hidden_states = hidden_states.narrow(0, token_offset, group_tokens) + group_probs = permuted_probs.reshape(-1).narrow(0, token_offset, group_tokens) + + # An all-empty expert group has no output rows and does not need weight + # quantization or a DeepGEMM launch. + if group_tokens == 0: + continue + + padded_group_tokens = sum( + ((count + _GROUPED_M_ALIGNMENT - 1) // _GROUPED_M_ALIGNMENT * _GROUPED_M_ALIGNMENT if count else 0) + for count in group_counts + ) + group_output = None + if direct_group_outputs: + assert output_storage is not None + group_output = output_storage.narrow( + 0, + token_offset, + padded_group_tokens, + ) + reuse_unpadded_input = reuse_input_buffer and padded_group_tokens == group_tokens + if reuse_unpadded_input: + # DeepEP normally supplies M-aligned expert counts. In the + # checkpoint/no-grad forward, activation quantization is the last + # reader of this group input, so the same slice can become FC1 + # gate/up scratch, SiLU/down scratch, and finally FC2 output. + group_output = group_hidden_states + + # In checkpoint recomputation the final-output slice is still unused + # here. Use it as the aligned BF16 expert input, quantize from it, and + # let FC2 overwrite it later. This removes another group-sized + # temporary (about 1.9 GiB for the largest EP32/8K group). + if reuse_unpadded_input: + padded_hidden_states = group_hidden_states + padded_counts = group_counts + valid_rows = None + else: + padded_hidden_states, padded_counts, valid_rows = _pad_expert_rows( + group_hidden_states, + group_counts, + output=group_output, + ) + if reuse_input_buffer: + # The padded expert-major input is necessary for DeepGEMM, but + # it becomes dead as soon as activation quantization finishes. + # Keep its storage as the FC1/activation/FC2 scratch instead + # of allocating those tensors alongside it. FC2 is compacted + # back into the original dispatched-input slice below. + group_output = padded_hidden_states + padded_tokens = padded_hidden_states.shape[0] + m_indices = _build_m_indices(padded_counts, device=hidden_states.device) + if m_indices.numel() != padded_tokens or padded_tokens % _GROUPED_M_ALIGNMENT: + raise RuntimeError( + "DeepGEMM contiguous grouped layout is not M-aligned: " + f"rows={padded_tokens}, indices={m_indices.numel()}, " + f"alignment={_GROUPED_M_ALIGNMENT}" + ) + + hidden_q, hidden_scale = _quantize_activation(padded_hidden_states, ops) + gate_up_shape = (padded_tokens, 2 * layout.ffn_hidden_size) + down_input_shape = (padded_tokens, layout.ffn_hidden_size) + gate_up = None + down_input = None + if group_output is not None: + # The padded expert input above has already been quantized, so its + # final-output destination is dead until FC2 writes the result. + # GLM's [hidden=6144, ffn=2048] layout fits gate/up and the SiLU + # result in two disjoint slices of that same BF16 storage. This + # removes the 628 MiB gate/up plus 314 MiB down-input allocations + # at the EP32/8K checkpoint-recompute peak. + gate_up_elements = math.prod(gate_up_shape) + down_input_elements = math.prod(down_input_shape) + if gate_up_elements + down_input_elements <= group_output.numel(): + output_scratch = group_output.reshape(-1) + gate_up = output_scratch.narrow(0, 0, gate_up_elements).view(gate_up_shape) + down_input = output_scratch.narrow( + 0, + gate_up_elements, + down_input_elements, + ).view(down_input_shape) + # Activation quantization has consumed padded_hidden_states. Until + # FC1 completes, the future SiLU/down-input slice is dead storage and + # is disjoint from FC1's gate/up output. Hold the grouped FP8 FC1 + # weights there instead of allocating 96 MiB at the full-pipeline + # checkpoint-recompute peak; stream ordering makes the slice reusable + # by silu_and_mul immediately after grouped_gemm returns. + fc1_qweight, fc1_scale = _quantize_grouped_weights( + module.linear_fc1, + expected_shape=layout.fc1_weight_shape, + layout=layout, + input_device=hidden_states.device, + module_name=f"{module_name}.linear_fc1", + ops=ops, + expert_start=expert_start, + expert_end=expert_end, + grouped_qweight_workspace=down_input, + ) + if gate_up is None and not reuse_input_buffer: + gate_up = _moe_recompute_scratch_view(module, hidden_states, gate_up_shape) + if gate_up is None: + gate_up = hidden_states.new_empty(gate_up_shape) + ops.grouped_gemm( + (hidden_q, hidden_scale), + (fc1_qweight, fc1_scale), + gate_up, + m_indices, + ) + del hidden_q, hidden_scale, fc1_qweight, fc1_scale + + if down_input is None: + down_input = hidden_states.new_empty(down_input_shape) + ops.silu_and_mul(gate_up, down_input) + del gate_up + + down_q, down_scale = _quantize_activation(down_input, ops) + del down_input + fc2_qweight, fc2_scale = _quantize_grouped_weights( + module.linear_fc2, + expected_shape=layout.fc2_weight_shape, + layout=layout, + input_device=hidden_states.device, + module_name=f"{module_name}.linear_fc2", + ops=ops, + expert_start=expert_start, + expert_end=expert_end, + ) + group_output_shape = (padded_tokens, layout.hidden_size) + if group_output is None and not reuse_input_buffer: + group_output = ( + _moe_recompute_scratch_view( + module, + hidden_states, + group_output_shape, + ) + if not reuse_input_buffer + else None + ) + if group_output is None: + group_output = hidden_states.new_empty(group_output_shape) + ops.grouped_gemm( + (down_q, down_scale), + (fc2_qweight, fc2_scale), + group_output, + m_indices, + ) + del down_q, down_scale, fc2_qweight, fc2_scale + + if valid_rows is not None: + group_output = _compact_valid_rows_inplace(group_output, valid_rows) + if not defer_router_probabilities: + group_output = _apply_router_probability_fp32_inplace(group_output, group_probs) + if not direct_group_outputs: + output_slice = output.narrow(0, token_offset, group_tokens) + if output_slice.data_ptr() != group_output.data_ptr(): + output_slice.copy_(group_output) + token_offset += group_tokens + + if token_offset != num_tokens: + raise AssertionError(f"Grouped MoE output row mismatch: {token_offset} != {num_tokens}") + return output + + +def _wrap_te_grouped_mlp( + module: torch.nn.Module, + module_name: str, +) -> bool: + if getattr(module, "_vime_deepgemm_moe_forward_wrapped", False): + return False + + _validate_parallelism() + layout = _validate_te_grouped_mlp(module, module_name) + fc1_weights = _get_expert_weights( + module.linear_fc1, + num_local_experts=layout.num_local_experts, + expected_shape=layout.fc1_weight_shape, + module_name=f"{module_name}.linear_fc1", + ) + fc2_weights = _get_expert_weights( + module.linear_fc2, + num_local_experts=layout.num_local_experts, + expected_shape=layout.fc2_weight_shape, + module_name=f"{module_name}.linear_fc2", + ) + if getattr(getattr(module, "config", None), "delay_wgrad_compute", False): + raise RuntimeError( + "DeepGEMM MoE custom backward does not support Megatron delay_wgrad_compute; " + "disable delay_wgrad_compute." + ) + + def deepgemm_moe_forward( + self, + permuted_local_hidden_states: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + ): + grad_enabled = torch.is_grad_enabled() + combine_buffer = getattr( + permuted_local_hidden_states, + _PREALLOCATED_COMBINE_BUFFER_ATTR, + None, + ) + if grad_enabled: + deepgemm_output = _DeepGEMMMoEWithBF16Backward.apply( + permuted_local_hidden_states, + tokens_per_expert, + permuted_probs, + self, + layout, + module_name, + *fc1_weights, + *fc2_weights, + ) + else: + deepgemm_output = _deepgemm_grouped_moe_forward( + self, + permuted_local_hidden_states, + tokens_per_expert, + permuted_probs, + layout=layout, + module_name=module_name, + reuse_input_buffer=True, + ) + if combine_buffer is not None: + setattr(deepgemm_output, _PREALLOCATED_COMBINE_BUFFER_ATTR, combine_buffer) + return deepgemm_output, getattr(self, "output_bias", None) + + module.forward = types.MethodType(deepgemm_moe_forward, module) + module._vime_deepgemm_moe_forward_wrapped = True + module._vime_deepgemm_moe_module_name = module_name + module._vime_deepgemm_moe_layout = layout + return True + + +def _get_global_layer_index(model_chunk: torch.nn.Module, module_name: str) -> int | None: + match = _LAYER_PATH_RE.search(module_name) + if match is None: + return None + layer_path = match.group("layer_path") + layer = model_chunk.get_submodule(layer_path) + layer_number = getattr(layer, "layer_number", None) + if layer_number is None: + raise RuntimeError( + "DeepGEMM MoE alignment requires TransformerLayer.layer_number to select " + f"global pipeline layers, but {layer_path} has no layer_number" + ) + return int(layer_number) - 1 + + +def _normalize_model_chunks(model) -> list[torch.nn.Module]: + if isinstance(model, torch.nn.Module): + return [model] + chunks = list(model) + if not all(isinstance(chunk, torch.nn.Module) for chunk in chunks): + raise RuntimeError("model must be a Megatron module or an iterable of model chunks") + return chunks + + +class _DeepEPScatterWithDeterministicBackward(torch.autograd.Function): + """Scatter routes with a deterministic backward.""" + + @staticmethod + def forward( + ctx, + hidden_states: torch.Tensor, + output_index: torch.Tensor, + total_rows: int, + ) -> torch.Tensor: + output = hidden_states.new_zeros((int(total_rows), hidden_states.shape[1])) + if hidden_states.is_cuda: + from vime.backends.megatron_utils.alignment.deterministic_route_kernels import scatter_routes_forward + + scatter_routes_forward( + hidden_states.contiguous(), + output_index.contiguous(), + output, + ) + else: + valid_positions = torch.nonzero(output_index >= 0, as_tuple=False) + for start in range(0, valid_positions.shape[0], _BACKWARD_CHUNK_ROWS): + positions = valid_positions[start : start + _BACKWARD_CHUNK_ROWS] + token_rows = positions[:, 0] + topk_columns = positions[:, 1] + destination_rows = output_index[token_rows, topk_columns].to(dtype=torch.long) + output.index_copy_( + 0, + destination_rows, + hidden_states.index_select(0, token_rows), + ) + + ctx.input_shape = tuple(hidden_states.shape) + ctx.input_dtype = hidden_states.dtype + ctx.input_device = hidden_states.device + ctx.save_for_backward(output_index) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + if not ctx.needs_input_grad[0]: + return None, None, None + + (output_index,) = ctx.saved_tensors + grad_input = torch.zeros( + ctx.input_shape, + dtype=ctx.input_dtype, + device=ctx.input_device, + ) + if grad_output.is_cuda: + from vime.backends.megatron_utils.alignment.deterministic_route_kernels import scatter_routes_backward + + scatter_routes_backward( + grad_output.contiguous(), + output_index.contiguous(), + grad_input, + ) + else: + # Accumulate top-k slots in their original order. Each token row + # is owned by one thread here, avoiding the nondeterministic + # atomics that an index_add over expert-major occurrences would + # require. + for start in range(0, output_index.shape[0], _BACKWARD_CHUNK_ROWS): + end = min(start + _BACKWARD_CHUNK_ROWS, output_index.shape[0]) + grad_chunk = grad_input[start:end] + for column in range(output_index.shape[1]): + route_rows = output_index[start:end, column].to(dtype=torch.long) + valid_rows = route_rows >= 0 + safe_route_rows = route_rows.clamp(min=0, max=max(grad_output.shape[0] - 1, 0)) + if grad_output.shape[0] == 0: + continue + selected = grad_output.index_select(0, safe_route_rows) + selected.masked_fill_(~valid_rows.unsqueeze(1), 0) + grad_chunk.add_(selected) + return grad_input, None, None + + +def _scatter_deepep_routes_with_padding( + hidden_states: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + return_route_positions: bool = False, + expected_route_count: int | None = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + bool, +]: + """Build VLLM's padded expert-major DeepEP layout deterministically.""" + if hidden_states.ndim != 2 or topk_indices.ndim != 2: + raise ValueError("DeepEP scatter expects [tokens, hidden] and [tokens, topk]") + if topk_indices.shape != topk_weights.shape: + raise ValueError("DeepEP top-k indices and weights must have identical shapes") + if hidden_states.shape[0] != topk_indices.shape[0]: + raise ValueError("DeepEP hidden-state and top-k token dimensions differ") + if topk_indices.dtype not in (torch.int32, torch.int64): + raise TypeError(f"DeepEP top-k indices must be integer, got {topk_indices.dtype}") + if topk_weights.dtype != torch.float32: + raise TypeError(f"DeepEP top-k weights must be FP32, got {topk_weights.dtype}") + + if tokens_per_expert.device.type == "cpu": + # Normal DeepEP returns this metadata as a CPU tensor constructed from + # its receive-count list. The output row count is consequently + # already known to Python; do not upload the counts only to block on a + # device sum immediately afterwards. + count_values = tuple(int(value) for value in tokens_per_expert.reshape(-1).tolist()) + total_rows = sum(count_values) + else: + count_values = None + + num_experts = tokens_per_expert.numel() + valid = (topk_indices >= 0) & (topk_indices < num_experts) + sanitized_indices = topk_indices.masked_fill(~valid, -1) + real_counts = torch.bincount( + sanitized_indices.masked_select(valid).to(dtype=torch.long), + minlength=num_experts, + ) + + # The route-preserving metadata handle gives Python the exact number of + # received routes. Normal DeepEP's CPU count list is unpadded in this + # case, so ``real_counts`` is the same per-expert metadata already resident + # on CUDA. Reusing it avoids a blocking pageable-CPU-to-CUDA upload once + # per MoE layer (and again during checkpoint recomputation). Retain the + # original upload for aligned/padded layouts and callers without the exact + # route count. + counts_are_exact_unpadded = ( + count_values is not None and expected_route_count is not None and total_rows == expected_route_count + ) + if counts_are_exact_unpadded: + counts = real_counts + torch._assert_async( + real_counts.sum() == expected_route_count, + "DeepEP received route count differs from its route-preserving metadata", + ) + else: + counts = tokens_per_expert.to( + device=topk_indices.device, + dtype=torch.long, + ).reshape(-1) + torch._assert_async( + torch.all(real_counts <= counts), + "DeepEP real route count exceeds its aligned expert count", + ) + + if count_values is None: + total_rows = int(counts.sum().item()) + permuted_probs = topk_weights.new_zeros((total_rows,)) + output_index = torch.full_like(topk_indices, -1) + routing_map = torch.zeros( + (topk_indices.shape[0], num_experts), + device=topk_indices.device, + dtype=torch.bool, + ) + expert_offsets = torch.cumsum(counts, dim=0) - counts + + if expected_route_count is not None and valid.is_cuda: + from vime.backends.megatron_utils.alignment.deterministic_route_kernels import compact_route_positions + + occurrences = compact_route_positions(valid.contiguous(), expected_route_count) + else: + occurrences = torch.nonzero(valid, as_tuple=False) + # The metadata handle exposes the compact route count as tensor shape. If + # it is available, use that static value instead of forcing nonzero to + # report its data-dependent output size to Python. + route_count = occurrences.shape[0] if expected_route_count is None else expected_route_count + all_routes_valid = route_count == topk_indices.numel() + if occurrences.numel(): + token_rows = occurrences[:, 0] + topk_columns = occurrences[:, 1] + route_experts = sanitized_indices[token_rows, topk_columns].to(dtype=torch.long) + expert_order = torch.argsort(route_experts, stable=True) + token_rows = token_rows.index_select(0, expert_order) + topk_columns = topk_columns.index_select(0, expert_order) + route_experts = route_experts.index_select(0, expert_order) + + real_offsets = torch.cumsum(real_counts, dim=0) - real_counts + within_expert = torch.arange( + route_experts.numel(), + device=hidden_states.device, + dtype=torch.long, + ) - real_offsets.index_select(0, route_experts) + destination_rows = expert_offsets.index_select(0, route_experts) + within_expert + permuted_probs.index_copy_( + 0, + destination_rows, + topk_weights[token_rows, topk_columns], + ) + output_index[token_rows, topk_columns] = destination_rows.to(dtype=output_index.dtype) + routing_map[token_rows, route_experts] = True + + torch._assert_async( + torch.all(~valid | (output_index >= 0)), + "A valid DeepEP route has no expert-major output row", + ) + permuted_hidden = _DeepEPScatterWithDeterministicBackward.apply( + hidden_states, + output_index, + total_rows, + ) + result = ( + permuted_hidden, + permuted_probs, + output_index, + sanitized_indices, + routing_map, + all_routes_valid, + ) + if return_route_positions: + return (*result, occurrences) + return result + + +class _VLLMEPGatherWithBF16Backward(torch.autograd.Function): + """VLLM's ordered FP32 gather with a deterministic BF16 backward.""" + + @staticmethod + def forward( + ctx, + hidden_states: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + output_index: torch.Tensor, + reuse_input_for_grad: bool, + static_mapping_valid: bool | None, + ) -> torch.Tensor: + if hidden_states.ndim != 2 or hidden_states.dtype != torch.bfloat16: + raise TypeError( + "VLLM ep_gather requires BF16 [expert_rows, hidden], got " + f"{hidden_states.dtype} {tuple(hidden_states.shape)}" + ) + if topk_indices.shape != topk_weights.shape or topk_indices.shape != output_index.shape: + raise ValueError("DeepEP gather IDs, weights, and output indices must align") + from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ep_gather + + output_shape = (topk_indices.shape[0], hidden_states.shape[1]) + output = torch.empty( + output_shape, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + ep_gather( + hidden_states, + topk_indices, + topk_weights, + output_index, + None, + output, + ) + ctx.reuse_input_for_grad = bool(reuse_input_for_grad) + ctx.static_mapping_valid = static_mapping_valid + ctx.save_for_backward(hidden_states, topk_weights, output_index) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden_states, topk_weights, output_index = ctx.saved_tensors + needs_hidden = ctx.needs_input_grad[0] + needs_weights = ctx.needs_input_grad[2] + if needs_hidden and ctx.reuse_input_for_grad: + # The caller guarantees this combine input has no forward consumer + # after ep_gather. Returning its detached storage as the input + # gradient avoids a second route-sized BF16 allocation. + grad_hidden = hidden_states.detach() + else: + grad_hidden = torch.zeros_like(hidden_states) if needs_hidden else None + grad_weights = torch.zeros_like(topk_weights) if needs_weights else None + + _ordered_route_backward( + route_values=hidden_states, + topk_weights=topk_weights, + output_index=output_index, + grad_output=grad_output, + grad_routes=grad_hidden, + grad_weights=grad_weights, + static_mapping_valid=ctx.static_mapping_valid, + ) + + return grad_hidden, None, grad_weights, None, None, None + + +def _compact_route_preserving_metadata_inputs( + hidden_states: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + *, + assume_all_routes_valid: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Build the small route-level dispatch used by the normal-mode prototype. + + The real hidden state remains rank-deduplicated in the primary DeepEP + dispatch. This second dispatch only carries a sixteen-BF16 fingerprint and + creates a normal DeepEP handle whose logical tokens are ``(token, slot)``. + The handle can consequently transport the real route outputs during + combine without multiplying dispatch hidden traffic by top-k. + + Sixteen BF16 values are the smallest payload accepted by DeepEP's normal + intranode dispatch: its TMA path requires the hidden payload to be a + multiple of 32 bytes. Keeping that constraint explicit here prevents a + device-side assertion for otherwise-valid route metadata. + """ + if hidden_states.ndim != 2 or hidden_states.dtype != torch.bfloat16: + raise TypeError( + "Route-preserving DeepEP metadata requires BF16 [tokens, hidden], " + f"got {hidden_states.dtype} {tuple(hidden_states.shape)}" + ) + if hidden_states.shape[1] < 16: + raise ValueError("Route-preserving DeepEP fingerprints require hidden >= 16") + if topk_indices.ndim != 2 or topk_weights.shape != topk_indices.shape: + raise ValueError("Route-preserving DeepEP top-k IDs and weights must align") + if topk_indices.shape[0] != hidden_states.shape[0]: + raise ValueError("Route-preserving DeepEP token counts do not align") + if topk_indices.dtype not in (torch.int32, torch.int64): + raise TypeError(f"Route-preserving DeepEP IDs must be integer, got {topk_indices.dtype}") + if topk_weights.dtype != torch.float32: + raise TypeError(f"Route-preserving DeepEP weights must be FP32, got {topk_weights.dtype}") + + flat_indices = topk_indices.reshape(-1) + if assume_all_routes_valid: + # With no capacity factor and no router token mask, the router emits a + # fixed valid top-k for every source token. Keep this fast path tied + # to that structural guarantee instead of synchronizing on nonzero's + # data-dependent output once per MoE layer. + torch._assert_async( + torch.all(flat_indices >= 0), + "Route-preserving DeepEP expected a valid fixed top-k", + ) + compact_indices = flat_indices.reshape(-1, 1).contiguous() + compact_weights = topk_weights.detach().reshape(-1, 1).contiguous() + fingerprints = ( + hidden_states.detach() + .narrow(1, 0, 16) + .unsqueeze(1) + .expand(-1, topk_indices.shape[1], -1) + .reshape(-1, 16) + .contiguous() + ) + output_index = torch.arange( + flat_indices.numel(), + device=flat_indices.device, + dtype=torch.long, + ).reshape_as(topk_indices) + return compact_indices, compact_weights, fingerprints, output_index, True + + valid_positions = torch.nonzero(flat_indices >= 0, as_tuple=False).reshape(-1) + if valid_positions.numel() == 0: + raise RuntimeError("Route-preserving DeepEP received no valid expert routes") + compact_indices = flat_indices.index_select(0, valid_positions).reshape(-1, 1).contiguous() + compact_weights = topk_weights.detach().reshape(-1).index_select(0, valid_positions).reshape(-1, 1).contiguous() + token_rows = torch.div(valid_positions, topk_indices.shape[1], rounding_mode="floor") + fingerprints = hidden_states.detach().narrow(1, 0, 16).index_select(0, token_rows).contiguous() + output_index = torch.full_like(topk_indices, -1, dtype=torch.long) + output_index.reshape(-1).index_copy_( + 0, + valid_positions, + torch.arange(valid_positions.numel(), device=valid_positions.device, dtype=torch.long), + ) + all_routes_valid = valid_positions.numel() == topk_indices.numel() + return compact_indices, compact_weights, fingerprints, output_index, all_routes_valid + + +def _deepep_route_handle_received_rows(handle: tuple) -> int: + """Return the received route count encoded by a normal DeepEP handle.""" + if not isinstance(handle, tuple): + raise TypeError(f"DeepEP route handle must be a tuple, got {type(handle).__name__}") + if len(handle) == 6: + # Intranode: (..., recv_src_idx, ...). + received_metadata = handle[3] + elif len(handle) == 10: + # Internode: (..., recv_src_meta, ...). + received_metadata = handle[7] + else: + raise ValueError(f"Unsupported normal DeepEP route handle length: {len(handle)}") + if not isinstance(received_metadata, torch.Tensor) or received_metadata.ndim < 1: + raise TypeError("DeepEP route handle has invalid received-source metadata") + return received_metadata.shape[0] + + +def _dispatch_route_preserving_deepep_metadata( + manager: object, + hidden_states: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + *, + assume_all_routes_valid: bool = False, +) -> tuple[tuple, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]: + """Create a route-level normal DeepEP handle with tiny payloads. + + This is intentionally a correctness prototype. Once validated, the same + route counts/source metadata are to be emitted by the primary aligned + dispatch so this extra metadata-only communication disappears. + """ + ( + route_indices, + route_weights, + route_fingerprints, + source_output_index, + all_routes_valid, + ) = _compact_route_preserving_metadata_inputs( + hidden_states, + topk_indices, + topk_weights, + assume_all_routes_valid=assume_all_routes_valid, + ) + from megatron.core.transformer.moe.fused_a2a import get_buffer, get_hidden_bytes + + group = manager.group + buffer = get_buffer(group, get_hidden_bytes(route_fingerprints)) + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + layout_event, + ) = buffer.get_dispatch_layout( + route_indices, + int(manager.num_experts), + async_finish=False, + allocate_on_comm_stream=False, + ) + ( + recv_fingerprints, + recv_route_indices, + recv_route_weights, + _, + route_handle, + _, + ) = buffer.dispatch( + route_fingerprints, + topk_idx=route_indices, + topk_weights=route_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=layout_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + if recv_route_indices is None or recv_route_weights is None: + raise RuntimeError("Route-preserving DeepEP metadata dispatch dropped top-k metadata") + return ( + route_handle, + recv_fingerprints, + recv_route_indices.reshape(-1), + recv_route_weights.reshape(-1), + source_output_index, + all_routes_valid, + ) + + +def _validate_and_order_route_preserving_outputs( + expert_outputs: torch.Tensor, + received_tokens: torch.Tensor, + received_topk_indices: torch.Tensor, + received_topk_weights: torch.Tensor, + output_index: torch.Tensor, + route_fingerprints: torch.Tensor, + route_indices: torch.Tensor, + route_weights: torch.Tensor, + *, + order_outputs: bool = True, + route_positions: torch.Tensor | None = None, +) -> torch.Tensor: + """Return expert outputs in the route handle's receive order. + + DeepEP currently produces the same source-token/slot order for the primary + rank-deduplicated dispatch and the virtual-token metadata dispatch. Do not + merely assume that invariant: validate expert ID, exact FP32 weight and an + sixteen-BF16 source fingerprint before using the route handle. + """ + if expert_outputs.ndim != 2 or received_tokens.ndim != 2: + raise ValueError("Route-preserving DeepEP expects 2D hidden tensors") + if received_topk_indices.shape != received_topk_weights.shape: + raise ValueError("Received DeepEP IDs and weights do not align") + if output_index.shape != received_topk_indices.shape: + raise ValueError("Received DeepEP route mapping does not align") + + positions = torch.nonzero(output_index >= 0, as_tuple=False) if route_positions is None else route_positions + if positions.shape[0] != route_indices.numel(): + raise RuntimeError( + "Route-preserving DeepEP route count mismatch: " + f"primary={positions.shape[0]} metadata={route_indices.numel()}" + ) + token_rows = positions[:, 0] + topk_slots = positions[:, 1] + expected_indices = received_topk_indices[token_rows, topk_slots].reshape(-1) + expected_weights = received_topk_weights[token_rows, topk_slots].reshape(-1) + expected_fingerprints = received_tokens.narrow(1, 0, 16).index_select(0, token_rows) + if route_fingerprints.shape != expected_fingerprints.shape: + raise RuntimeError( + "Route-preserving DeepEP fingerprint shape mismatch: " + f"{tuple(route_fingerprints.shape)} != {tuple(expected_fingerprints.shape)}" + ) + + torch._assert_async( + torch.all(expected_indices == route_indices.to(dtype=expected_indices.dtype)), + "Route-preserving DeepEP metadata changed local expert order", + ) + torch._assert_async( + torch.all(expected_weights == route_weights.to(dtype=expected_weights.dtype)), + "Route-preserving DeepEP metadata changed route probability order", + ) + torch._assert_async( + torch.all(expected_fingerprints == route_fingerprints), + "Route-preserving DeepEP metadata changed source-token order", + ) + + if not order_outputs: + return expert_outputs + route_rows = output_index[token_rows, topk_slots].to(dtype=torch.long) + return expert_outputs.index_select(0, route_rows) + + +def _patch_vllm_deepep_layer(mlp: torch.nn.Module, global_layer: int) -> bool: + """Match VLLM low-latency reduction over Megatron normal DeepEP.""" + if getattr(mlp, "_vime_vllm_deepep_alignment", False): + return False + + router = getattr(mlp, "router", None) + dispatcher = getattr(mlp, "token_dispatcher", None) + experts = getattr(mlp, "experts", None) + if router is None or dispatcher is None or experts is None: + raise RuntimeError(f"Layer {global_layer} is not a complete MoE layer") + manager = getattr(dispatcher, "_comm_manager", None) + if manager is None or manager.__class__.__name__ != "_DeepepManager": + raise RuntimeError( + "VLLM DeepEP alignment requires MCore's flex _DeepepManager; " + f"layer {global_layer} has {type(manager).__name__}" + ) + scaling_factor = float(router.config.moe_router_topk_scaling_factor or 1.0) + original_routing = router.routing + + def routing_without_final_scaling( + patched_router: torch.nn.Module, + *args: object, + **kwargs: object, + ): + config = patched_router.config + previous = config.moe_router_topk_scaling_factor + config.moe_router_topk_scaling_factor = 1.0 + try: + return original_routing(*args, **kwargs) + finally: + config.moe_router_topk_scaling_factor = previous + + router.routing = types.MethodType(routing_without_final_scaling, router) + original_setup_metadata = manager.setup_metadata + original_dispatch = manager.dispatch + from vime.utils.routing_replay import consume_ordered_topk, register_ordered_topk_capture + + register_ordered_topk_capture(router) + + def setup_ordered_metadata( + patched_manager, + routing_map: torch.Tensor, + probs: torch.Tensor, + router_token_masks: torch.Tensor | None = None, + ) -> None: + patched_manager._vime_source_fixed_topk_valid = False + ordered = consume_ordered_topk(router) + if ordered is None: + # Megatron 1dcf0dafa's DeepEP dispatcher setup_metadata takes only + # (routing_map, probs); the router_token_masks padding-mask arg exists + # on newer forks. Call the original with whatever arity it supports. + if router_token_masks is None: + original_setup_metadata(routing_map, probs) + else: + original_setup_metadata(routing_map, probs, router_token_masks) + return + + num_tokens = routing_map.shape[0] + if ordered.shape[0] != num_tokens: + raise ValueError( + "Ordered top-k token count differs from DeepEP input: " f"{ordered.shape[0]} != {num_tokens}" + ) + ordered = ordered.to( + device=probs.device, + dtype=torch.int64, + non_blocking=ordered.is_pinned(), + ) + dense_probs = probs.reshape(num_tokens, patched_manager.num_experts) + patched_manager.token_indices = ordered + patched_manager.token_probs = dense_probs.gather(-1, ordered) + patched_manager._vime_source_fixed_topk_valid = ( + patched_manager.capacity_factor is None and router_token_masks is None + ) + if patched_manager.capacity_factor is not None: + patched_manager.token_indices = patched_manager.token_indices.masked_fill( + patched_manager.token_probs == 0, + -1, + ) + if router_token_masks is not None: + patched_manager.token_indices = patched_manager.token_indices.masked_fill( + router_token_masks.view(-1, 1), + -1, + ) + + manager.setup_metadata = types.MethodType( + setup_ordered_metadata, + manager, + ) + + def dispatch_with_route_preserving_handle( + patched_manager, + hidden_states: torch.Tensor, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + ) -> torch.Tensor: + if patched_manager.token_indices is None or patched_manager.token_probs is None: + raise RuntimeError("Ordered source top-k metadata is unavailable before DeepEP dispatch") + dispatched_hidden = original_dispatch( + hidden_states, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + # _DeepepManager normalizes probabilities to FP32 in its dispatch + # preamble. Capture the normalized tensors that were actually sent. + source_topk_indices = patched_manager.token_indices + source_topk_weights = patched_manager.token_probs + source_fixed_topk_valid = bool(getattr(patched_manager, "_vime_source_fixed_topk_valid", False)) + ( + route_handle, + recv_route_fingerprints, + recv_route_indices, + recv_route_weights, + source_output_index, + source_all_routes_valid, + ) = _dispatch_route_preserving_deepep_metadata( + patched_manager, + hidden_states, + source_topk_indices, + source_topk_weights, + assume_all_routes_valid=source_fixed_topk_valid, + ) + route_metadata_prevalidated = False + patched_manager._vime_route_handle = route_handle + patched_manager._vime_route_recv_fingerprints = recv_route_fingerprints + patched_manager._vime_route_recv_indices = recv_route_indices + patched_manager._vime_route_recv_weights = recv_route_weights + patched_manager._vime_route_metadata_prevalidated = route_metadata_prevalidated + patched_manager._vime_route_source_topk_indices = source_topk_indices + patched_manager._vime_route_source_topk_weights = source_topk_weights + patched_manager._vime_route_source_output_index = source_output_index + patched_manager._vime_route_source_all_valid = source_all_routes_valid + return dispatched_hidden + + manager.dispatch = types.MethodType( + dispatch_with_route_preserving_handle, + manager, + ) + + def get_permuted_hidden_states_by_experts( + patched_manager, + hidden_states: torch.Tensor, + ): + topk_indices = patched_manager.dispatched_indices + topk_weights = patched_manager.dispatched_probs + if topk_indices is None or topk_weights is None: + raise RuntimeError("DeepEP dispatch metadata is unavailable before local permutation") + ( + permuted_hidden, + permuted_probs, + output_index, + sanitized_indices, + routing_map, + all_routes_valid, + route_positions, + ) = _scatter_deepep_routes_with_padding( + hidden_states, + topk_indices, + topk_weights, + patched_manager.tokens_per_expert, + return_route_positions=True, + expected_route_count=_deepep_route_handle_received_rows(patched_manager._vime_route_handle), + ) + patched_manager.hidden_shape_before_permute = hidden_states.shape + patched_manager.dispatched_routing_map = routing_map + patched_manager._vime_vllm_topk_indices = sanitized_indices + patched_manager._vime_vllm_topk_weights = topk_weights + patched_manager._vime_vllm_output_index = output_index + patched_manager._vime_vllm_route_positions = route_positions + patched_manager._vime_vllm_all_routes_valid = all_routes_valid + patched_manager._vime_vllm_expert_inputs = permuted_hidden + patched_manager._vime_vllm_expert_probs = permuted_probs + patched_manager._vime_vllm_tokens_per_expert = patched_manager.tokens_per_expert + route_fingerprints = getattr( + patched_manager, + "_vime_route_recv_fingerprints", + None, + ) + route_indices = getattr(patched_manager, "_vime_route_recv_indices", None) + route_weights = getattr(patched_manager, "_vime_route_recv_weights", None) + route_metadata_prevalidated = bool(getattr(patched_manager, "_vime_route_metadata_prevalidated", False)) + if route_metadata_prevalidated: + if any(value is not None for value in (route_fingerprints, route_indices, route_weights)): + raise RuntimeError("Cached DeepEP route metadata retained unexpected payloads") + else: + if route_fingerprints is None or route_indices is None or route_weights is None: + raise RuntimeError("Route-preserving DeepEP metadata handle is unavailable") + _validate_and_order_route_preserving_outputs( + permuted_hidden, + hidden_states, + sanitized_indices, + topk_weights, + output_index, + route_fingerprints, + route_indices, + route_weights, + order_outputs=False, + route_positions=route_positions, + ) + del patched_manager._vime_route_recv_fingerprints + del patched_manager._vime_route_recv_indices + del patched_manager._vime_route_recv_weights + del patched_manager._vime_route_metadata_prevalidated + return permuted_hidden, permuted_probs + + def get_restored_hidden_states_by_experts( + patched_manager, + hidden_states: torch.Tensor, + ): + topk_indices = getattr(patched_manager, "_vime_vllm_topk_indices", None) + topk_weights = getattr(patched_manager, "_vime_vllm_topk_weights", None) + output_index = getattr(patched_manager, "_vime_vllm_output_index", None) + route_positions = getattr(patched_manager, "_vime_vllm_route_positions", None) + if topk_indices is None or topk_weights is None or output_index is None or route_positions is None: + raise RuntimeError("Saved route-preserving DeepEP mapping is unavailable") + token_rows = route_positions[:, 0] + topk_slots = route_positions[:, 1] + route_rows = output_index[token_rows, topk_slots].to(dtype=torch.long) + output = hidden_states.index_select(0, route_rows) + del patched_manager._vime_vllm_topk_indices + del patched_manager._vime_vllm_topk_weights + del patched_manager._vime_vllm_output_index + del patched_manager._vime_vllm_route_positions + del patched_manager._vime_vllm_all_routes_valid + del patched_manager._vime_vllm_expert_inputs + del patched_manager._vime_vllm_expert_probs + del patched_manager._vime_vllm_tokens_per_expert + return output + + manager.get_permuted_hidden_states_by_experts = types.MethodType( + get_permuted_hidden_states_by_experts, + manager, + ) + manager.get_restored_hidden_states_by_experts = types.MethodType( + get_restored_hidden_states_by_experts, + manager, + ) + + # Megatron 1dcf0dafa's MoELayer.combine(output) runs only token_combine, with a + # separate postprocess(output, shared_expert_output) doing combine_postprocess + + # shared-expert add. Newer forks fuse all of that into + # combine(output, shared_expert_output). Detect which convention applies here. + import inspect + + combine_fuses_postprocess = len(inspect.signature(mlp.combine).parameters) >= 2 + + if not combine_fuses_postprocess: + # Megatron 1dcf0dafa split token-combine from postprocess. Preserve + # VLLM's single BF16 shared.add_(routed, alpha=scaling_factor) + # operation here: scaling the routed value first would introduce an + # extra BF16 rounding before the shared-expert addition. + def postprocess( + patched_mlp: torch.nn.Module, + output: torch.Tensor, + shared_expert_output: torch.Tensor | None, + ) -> torch.Tensor: + output = patched_mlp.token_dispatcher.combine_postprocess(output) + if bool(getattr(patched_mlp.config, "moe_latent_size", None)): + output, _ = patched_mlp.fc2_latent_proj(output) + if shared_expert_output is not None: + return torch.add( + shared_expert_output, + output, + alpha=scaling_factor, + ) + if scaling_factor != 1.0: + output = output * scaling_factor + return output + + mlp.postprocess = types.MethodType(postprocess, mlp) + + def combine( + patched_mlp: torch.nn.Module, + output: torch.Tensor, + shared_expert_output: torch.Tensor | None = None, + ) -> torch.Tensor: + route_handle = getattr(manager, "_vime_route_handle", None) + source_topk_indices = getattr(manager, "_vime_route_source_topk_indices", None) + source_topk_weights = getattr(manager, "_vime_route_source_topk_weights", None) + source_output_index = getattr(manager, "_vime_route_source_output_index", None) + source_all_routes_valid = getattr(manager, "_vime_route_source_all_valid", None) + if ( + route_handle is None + or source_topk_indices is None + or source_topk_weights is None + or source_output_index is None + or source_all_routes_valid is None + ): + raise RuntimeError("Route-preserving DeepEP combine handle is incomplete") + from megatron.core.transformer.moe.fused_a2a import fused_combine + + combined_routes, _ = fused_combine( + output, + manager.group, + route_handle, + async_finish=True, + allocate_on_comm_stream=getattr(patched_mlp.token_dispatcher, "allocate_on_comm_stream", False), + ) + output = _VLLMEPGatherWithBF16Backward.apply( + combined_routes, + source_topk_indices, + source_topk_weights, + source_output_index, + True, + source_all_routes_valid, + ) + manager.handle = None + del manager._vime_route_handle + del manager._vime_route_source_topk_indices + del manager._vime_route_source_topk_weights + del manager._vime_route_source_output_index + del manager._vime_route_source_all_valid + if combine_fuses_postprocess: + output = patched_mlp.token_dispatcher.combine_postprocess(output) + if shared_expert_output is not None: + output = torch.add( + shared_expert_output, + output, + alpha=scaling_factor, + ) + elif scaling_factor != 1.0: + output = output * scaling_factor + # 1dcf0dafa applies reshape, routed scaling, and the shared-expert add + # in the patched postprocess above so their rounding order stays fused. + return output + + mlp.combine = types.MethodType(combine, mlp) + experts._vime_defer_router_probabilities = True + experts._vime_reuse_expert_input_for_grad = True + mlp._vime_vllm_deepep_alignment = True + mlp._vime_vllm_deepep_global_layer = global_layer + return True + + +def enable_vllm_deepep_moe_alignment(args, model, store_prefix: str) -> None: + """Install the exact VLLM DeepEP combine semantics on selected MoE layers.""" + del store_prefix + if not bool(getattr(args, "deterministic_mode", False)): + return + if not bool(getattr(args, "moe_enable_deepep", False)): + return + selected_layers = {int(layer) for layer in (getattr(args, "megatron_deepgemm_moe_forward_layers", None) or ())} + if not selected_layers: + return + + patched = [] + for model_chunk in _normalize_model_chunks(model): + for layer_name, layer in model_chunk.named_modules(): + if re.match(r"^(?:.*\.)?decoder\.layers\.\d+$", layer_name) is None: + continue + layer_number = getattr(layer, "layer_number", None) + if layer_number is None: + continue + global_layer = int(layer_number) - 1 + if global_layer not in selected_layers: + continue + mlp = getattr(layer, "mlp", None) + if mlp is None or not hasattr(mlp, "experts"): + continue + if _patch_vllm_deepep_layer(mlp, global_layer): + patched.append(global_layer) + + if patched and _should_log_deepgemm_summary(): + logger.info( + "Enabled VLLM ordered FP32 DeepEP gather on %d local MoE layers " "(global layers=%s)", + len(patched), + _format_int_ranges(selected_layers), + ) + + +def install_deepgemm_moe_forward( + model, + global_layer_indices: Iterable[int], + *, + target_suffixes: Iterable[str] = _DEFAULT_TARGET_SUFFIXES, +) -> list[str]: + """Wrap selected global MoE layers, using their ``layer_number`` metadata. + + ``global_layer_indices`` are zero-based global decoder-layer indices. This + remains correct with pipeline parallelism even though each model chunk's + ``decoder.layers`` list starts at local index zero. + """ + _validate_parallelism() + selected_layers = {int(layer_index) for layer_index in global_layer_indices} + if not selected_layers: + raise RuntimeError("global_layer_indices must select at least one MoE layer") + suffixes = tuple(target_suffixes) + if not suffixes: + raise RuntimeError("target_suffixes must select at least one module name") + + workspace_bytes = _combine_workspace_bytes() + wrapped = [] + workspaces_by_device: dict[torch.device, torch.Tensor] = {} + for model_chunk in _normalize_model_chunks(model): + for name, module in model_chunk.named_modules(): + if not any(name.endswith(suffix) for suffix in suffixes): + continue + if _get_global_layer_index(model_chunk, name) not in selected_layers: + continue + if _wrap_te_grouped_mlp(module, name): + mlp_name = name.rsplit(".", 1)[0] + mlp = model_chunk.get_submodule(mlp_name) + dispatcher = getattr(mlp, "token_dispatcher", None) + if ( + workspace_bytes is not None + and dispatcher is not None + and int(getattr(dispatcher, "num_local_experts", 1)) > 1 + ): + _wrap_preallocated_combine_preprocess(dispatcher) + parameter = next(module.parameters()) + workspace = workspaces_by_device.get(parameter.device) + if workspace is None: + workspace = torch.empty( + workspace_bytes, + dtype=torch.uint8, + device=parameter.device, + ) + workspaces_by_device[parameter.device] = workspace + setattr(dispatcher, _COMBINE_WORKSPACE_ATTR, workspace) + setattr(module, _COMBINE_WORKSPACE_ATTR, workspace) + _wrap_preallocated_dispatch_postprocess(dispatcher) + _wrap_preallocated_token_combine(dispatcher) + wrapped.append(name) + + if wrapped and _should_log_deepgemm_summary(): + logger.info( + "Enabled VLLM grouped DeepGEMM MoE forward+BF16-backward on %d " "TEGroupedMLPs (global layers=%s)", + len(wrapped), + _format_int_ranges(selected_layers), + ) + logger.debug("DeepGEMM wrapped TEGroupedMLPs: %s", ", ".join(wrapped)) + else: + # Most PP ranks legitimately do not own a requested global layer. + logger.debug( + "No TEGroupedMLP matched the requested DeepGEMM MoE layers %s and suffixes %s", + sorted(selected_layers), + suffixes, + ) + return wrapped + + +def enable_deepgemm_moe_forward(args, model, store_prefix: str) -> None: + """Install the grouped DeepGEMM forward-value probe on selected MoE layers.""" + del store_prefix + layers = getattr(args, "megatron_deepgemm_moe_forward_layers", None) + if layers is None: + raise RuntimeError( + "args.megatron_deepgemm_moe_forward_layers is required; pass --megatron-deepgemm-moe-forward-layers" + ) + suffixes = getattr(args, "megatron_deepgemm_moe_forward_modules", None) or _DEFAULT_TARGET_SUFFIXES + install_deepgemm_moe_forward( + model, + layers, + target_suffixes=suffixes, + ) diff --git a/vime/backends/megatron_utils/alignment/deterministic_route_kernels.py b/vime/backends/megatron_utils/alignment/deterministic_route_kernels.py new file mode 100644 index 000000000..1bd2127a1 --- /dev/null +++ b/vime/backends/megatron_utils/alignment/deterministic_route_kernels.py @@ -0,0 +1,295 @@ +"""Deterministic Triton kernels for route permutation gradients. + +These kernels only replace launch-heavy tensor indexing with equivalent +pointwise copies or one-program-per-token ordered reductions. They do not use +atomics, and every visible output element has exactly one writer. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _scatter_routes_forward_kernel( + hidden_states, + output_index, + output, + num_slots, + hidden_size: tl.constexpr, + topk: tl.constexpr, + BLOCK_D: tl.constexpr, +): + hidden_block = tl.program_id(0) + first_slot = tl.program_id(1) + num_slot_programs = tl.num_programs(1) + hidden_offsets = hidden_block * BLOCK_D + tl.arange(0, BLOCK_D) + hidden_mask = hidden_offsets < hidden_size + + for slot_i32 in range(first_slot, num_slots, num_slot_programs): + slot = slot_i32.to(tl.int64) + destination_i32 = tl.load(output_index + slot) + destination = destination_i32.to(tl.int64) + token = slot // topk + valid = destination_i32 >= 0 + values = tl.load( + hidden_states + token * hidden_size + hidden_offsets, + mask=hidden_mask & valid, + other=0.0, + ) + tl.store( + output + destination * hidden_size + hidden_offsets, + values, + mask=hidden_mask & valid, + ) + + +@triton.jit +def _scatter_routes_backward_kernel( + grad_output, + output_index, + grad_input, + num_tokens, + grad_output_rows, + hidden_size: tl.constexpr, + topk: tl.constexpr, + BLOCK_D: tl.constexpr, +): + hidden_block = tl.program_id(0) + first_token = tl.program_id(1) + num_token_programs = tl.num_programs(1) + hidden_offsets = hidden_block * BLOCK_D + tl.arange(0, BLOCK_D) + hidden_mask = hidden_offsets < hidden_size + + for token_i32 in range(first_token, num_tokens, num_token_programs): + token = token_i32.to(tl.int64) + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) + for route_slot in range(topk): + route_i32 = tl.load(output_index + token * topk + route_slot) + valid = (route_i32 >= 0) & (route_i32 < grad_output_rows) + safe_route = tl.maximum(route_i32, 0).to(tl.int64) + route_grad = tl.load( + grad_output + safe_route * hidden_size + hidden_offsets, + mask=hidden_mask & valid, + other=0.0, + ).to(tl.float32) + # The reference performs one BF16 in-place add per top-k slot. + # Keep that observable rounding boundary instead of accumulating + # all slots in FP32 and casting only once. + accumulator = (accumulator + route_grad).to(tl.bfloat16).to(tl.float32) + tl.store( + grad_input + token * hidden_size + hidden_offsets, + accumulator, + mask=hidden_mask, + ) + + +@triton.jit +def _ordered_route_grad_kernel( + grad_output, + topk_weights, + output_index, + grad_routes, + num_slots, + hidden_size: tl.constexpr, + topk: tl.constexpr, + BLOCK_D: tl.constexpr, +): + hidden_block = tl.program_id(0) + first_slot = tl.program_id(1) + num_slot_programs = tl.num_programs(1) + hidden_offsets = hidden_block * BLOCK_D + tl.arange(0, BLOCK_D) + hidden_mask = hidden_offsets < hidden_size + + for slot_i32 in range(first_slot, num_slots, num_slot_programs): + slot = slot_i32.to(tl.int64) + route_i32 = tl.load(output_index + slot) + route = route_i32.to(tl.int64) + token = slot // topk + valid = route_i32 >= 0 + weight = tl.load(topk_weights + slot).to(tl.float32) + token_grad = tl.load( + grad_output + token * hidden_size + hidden_offsets, + mask=hidden_mask & valid, + other=0.0, + ).to(tl.float32) + tl.store( + grad_routes + route * hidden_size + hidden_offsets, + token_grad * weight, + mask=hidden_mask & valid, + ) + + +@triton.jit +def _compact_route_positions_kernel( + valid, + prefix, + positions, + num_slots, + topk: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + flat_slot = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + in_bounds = flat_slot < num_slots + is_valid = tl.load(valid + flat_slot, mask=in_bounds, other=0).to(tl.int1) + compact_row = tl.load(prefix + flat_slot, mask=in_bounds, other=0).to(tl.int64) - 1 + write_mask = in_bounds & is_valid + tl.store( + positions + compact_row * 2, + flat_slot // topk, + mask=write_mask, + ) + tl.store( + positions + compact_row * 2 + 1, + flat_slot % topk, + mask=write_mask, + ) + + +def _validate_common( + value: torch.Tensor, + output_index: torch.Tensor, +) -> None: + if not value.is_cuda or not output_index.is_cuda: + raise RuntimeError("deterministic route kernels require CUDA tensors") + if value.ndim != 2 or value.dtype != torch.bfloat16: + raise TypeError( + "deterministic route kernels require BF16 [rows, hidden], got " f"{value.dtype} {tuple(value.shape)}" + ) + if output_index.ndim != 2 or output_index.dtype not in (torch.int32, torch.int64): + raise TypeError( + "deterministic route kernels require an integer [tokens, topk] mapping, got " + f"{output_index.dtype} {tuple(output_index.shape)}" + ) + if not value.is_contiguous() or not output_index.is_contiguous(): + raise RuntimeError("deterministic route kernel inputs must be contiguous") + + +def scatter_routes_forward( + hidden_states: torch.Tensor, + output_index: torch.Tensor, + output: torch.Tensor, +) -> None: + """Copy token rows into their unique expert-major route rows.""" + _validate_common(hidden_states, output_index) + if output.ndim != 2 or output.dtype != hidden_states.dtype or not output.is_contiguous(): + raise TypeError("deterministic route-scatter output must be contiguous BF16 [rows, hidden]") + hidden_size = hidden_states.shape[1] + block_d = 1024 if hidden_size >= 1024 else triton.next_power_of_2(hidden_size) + num_slot_programs = min(output_index.numel(), 8192) + grid = (triton.cdiv(hidden_size, block_d), num_slot_programs) + _scatter_routes_forward_kernel[grid]( + hidden_states, + output_index, + output, + output_index.numel(), + hidden_size=hidden_size, + topk=output_index.shape[1], + BLOCK_D=block_d, + num_warps=4, + ) + + +def scatter_routes_backward( + grad_output: torch.Tensor, + output_index: torch.Tensor, + grad_input: torch.Tensor, +) -> None: + """Sum route gradients in fixed top-k order into token rows.""" + _validate_common(grad_output, output_index) + if ( + grad_input.shape != (output_index.shape[0], grad_output.shape[1]) + or grad_input.dtype != grad_output.dtype + or not grad_input.is_contiguous() + ): + raise TypeError("deterministic route-scatter input gradient has an invalid layout") + hidden_size = grad_output.shape[1] + block_d = 1024 if hidden_size >= 1024 else triton.next_power_of_2(hidden_size) + num_token_programs = min(output_index.shape[0], 2048) + grid = (triton.cdiv(hidden_size, block_d), num_token_programs) + _scatter_routes_backward_kernel[grid]( + grad_output, + output_index, + grad_input, + output_index.shape[0], + grad_output.shape[0], + hidden_size=hidden_size, + topk=output_index.shape[1], + BLOCK_D=block_d, + num_warps=4, + ) + + +def ordered_route_grad( + grad_output: torch.Tensor, + topk_weights: torch.Tensor, + output_index: torch.Tensor, + grad_routes: torch.Tensor, +) -> None: + """Write each token/slot gradient to its unique route row.""" + _validate_common(grad_output, output_index) + if topk_weights.shape != output_index.shape or topk_weights.dtype != torch.float32: + raise TypeError("ordered route gradients require FP32 [tokens, topk] weights") + if not topk_weights.is_cuda or not topk_weights.is_contiguous(): + raise RuntimeError("ordered route-gradient weights must be contiguous CUDA tensors") + if ( + grad_routes.ndim != 2 + or grad_routes.shape[1] != grad_output.shape[1] + or grad_routes.dtype != grad_output.dtype + or not grad_routes.is_cuda + or not grad_routes.is_contiguous() + ): + raise TypeError("ordered route-gradient output has an invalid layout") + hidden_size = grad_output.shape[1] + block_d = 1024 if hidden_size >= 1024 else triton.next_power_of_2(hidden_size) + num_slot_programs = min(output_index.numel(), 8192) + grid = (triton.cdiv(hidden_size, block_d), num_slot_programs) + _ordered_route_grad_kernel[grid]( + grad_output, + topk_weights, + output_index, + grad_routes, + output_index.numel(), + hidden_size=hidden_size, + topk=output_index.shape[1], + BLOCK_D=block_d, + num_warps=4, + ) + + +def compact_route_positions( + valid: torch.Tensor, + num_routes: int, +) -> torch.Tensor: + """Compact a row-major fixed-top-k validity mask without a host sync.""" + if valid.ndim != 2 or valid.dtype != torch.bool or not valid.is_cuda: + raise ValueError("route-position compaction requires a CUDA bool [tokens, topk] mask") + if num_routes < 0 or num_routes > valid.numel(): + raise ValueError(f"invalid compact route count {num_routes} for {valid.numel()} slots") + + flat_valid = valid.reshape(-1).contiguous() + prefix = torch.cumsum(flat_valid, dim=0, dtype=torch.int32) + positions = torch.empty( + (num_routes, 2), + device=valid.device, + dtype=torch.long, + ) + if flat_valid.numel(): + torch._assert_async( + prefix[-1] == num_routes, + "DeepEP compact route count differs from its metadata handle", + ) + if num_routes: + block_size = 256 + grid = (triton.cdiv(flat_valid.numel(), block_size),) + _compact_route_positions_kernel[grid]( + flat_valid, + prefix, + positions, + flat_valid.numel(), + topk=valid.shape[1], + BLOCK_SIZE=block_size, + ) + return positions diff --git a/vime/backends/megatron_utils/alignment/env.py b/vime/backends/megatron_utils/alignment/env.py new file mode 100644 index 000000000..3699d36da --- /dev/null +++ b/vime/backends/megatron_utils/alignment/env.py @@ -0,0 +1,39 @@ +"""Deterministic train/rollout alignment environment. + +Centralizes the numerical-alignment environment variables that both train +(Megatron) and rollout (VLLM) actors must share for GLM-5 train/rollout +log-prob alignment. Launchers (and the 6-layer gate test) merge this with +their own connectivity settings (``PYTHONPATH``, ``MASTER_ADDR``, NIC names, +proxy, IBGDA handler), which are cluster-specific and intentionally not here. +""" + +from __future__ import annotations + + +def alignment_env(*, kv_fp8_qat: bool = False) -> dict[str, str]: + """Return the shared deterministic-alignment env vars. + + ``kv_fp8_qat`` enables the FP8-E4M3 KV-cache QAT path (bf16 KV when False). + """ + return { + # Deterministic collectives / matmul. + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_P2P_LEVEL": "NVL", + "NCCL_ALGO": "^NVLS", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "TORCH_COMPILE_DISABLE": "1", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "TE_DISABLE_FA3": "TRUE", + "NVSHMEM_DISABLE_NCCL": "1", + # DeepGEMM batch-invariant FP8 forward. + "VLLM_BATCH_INVARIANT": "1", + # Megatron train side borrows VLLM's aligned kernels. + "MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS": "1", + "MEGATRON_USE_VLLM_FP8_INDEXER": "1", + "MEGATRON_USE_VLLM_ROUTER_GEMM": "1", + "MEGATRON_USE_VLLM_ROPE": "1", + "MEGATRON_USE_VLLM_SPARSE_MLA": "1", + # DSA KV cache dtype. + "DSA_KV_FP8_QAT": "1" if kv_fp8_qat else "0", + "DSA_KV_FP8_QAT_BLOCK_SIZE": "128", + } diff --git a/vime/backends/megatron_utils/alignment/layerwise_alignment.py b/vime/backends/megatron_utils/alignment/layerwise_alignment.py new file mode 100644 index 000000000..ba23cc1fc --- /dev/null +++ b/vime/backends/megatron_utils/alignment/layerwise_alignment.py @@ -0,0 +1,145 @@ +"""Opt-in Megatron layer-output dumps for train/rollout alignment gates.""" + +from __future__ import annotations + +import logging +import os +import re +from functools import partial +from pathlib import Path +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_LAYER_RE = re.compile(r"^(?:.*\.)?decoder\.layers\.(\d+)$") + + +def _global_rank() -> int: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + return 0 + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + if isinstance(value, dict): + for item in value.values(): + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + + +class _MegatronLayerwiseDumper: + def __init__(self, dump_dir: str, selected_layers: set[int], store_prefix: str): + self.dump_dir = Path(dump_dir) / f"rank{_global_rank():05d}" + self.dump_dir.mkdir(parents=True, exist_ok=True) + self.selected_layers = selected_layers + self.store_prefix = store_prefix.rstrip("_") or "actor" + self.module_suffixes = tuple( + suffix.strip() + for suffix in os.getenv("VIME_LAYERWISE_ALIGNMENT_MODULE_SUFFIXES", "").split(",") + if suffix.strip() + ) + self.pass_id = 0 + self.current: dict[str, Any] = {} + + def pre_forward(self, module, args, kwargs): + del module, args + self.current = { + "store_prefix": self.store_prefix, + "layers": {}, + "modules": {}, + } + input_ids = kwargs.get("input_ids") + if isinstance(input_ids, torch.Tensor): + self.current["input_ids"] = input_ids.detach().cpu() + packed_seq_params = kwargs.get("packed_seq_params") + cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None) + if isinstance(cu_seqlens, torch.Tensor): + self.current["cu_seqlens"] = cu_seqlens.detach().cpu() + + def record_layer(self, layer_id: int, module, args, output): + del module, args + tensor = _first_tensor(output) + if tensor is not None: + self.current.setdefault("layers", {})[layer_id] = tensor.detach().cpu() + + def record_module(self, module_name: str, module, args, output): + del module, args + tensor = _first_tensor(output) + if tensor is not None: + self.current.setdefault("modules", {})[module_name] = tensor.detach().cpu() + + def post_forward(self, module, args, output): + del module, args, output + if "input_ids" not in self.current: + raise RuntimeError("Megatron layerwise dump did not observe model input_ids") + observed_layers = set(self.current.get("layers", {})) + missing_layers = self.selected_layers - observed_layers + if missing_layers: + raise RuntimeError("Megatron layerwise dump missed selected layers: " f"{sorted(missing_layers)}") + output_path = self.dump_dir / f"{self.store_prefix}_Pass{self.pass_id:05d}.pt" + torch.save(self.current, output_path) + logger.info("Dumped Megatron layer outputs to %s", output_path) + self.pass_id += 1 + self.current = {} + + def register(self, model_chunk) -> int: + model_chunk.register_forward_pre_hook(self.pre_forward, with_kwargs=True) + model_chunk.register_forward_hook(self.post_forward) + registered_layers = 0 + for module_name, module in model_chunk.named_modules(): + match = _LAYER_RE.match(module_name) + if match is None: + continue + layer_number = getattr(module, "layer_number", None) + layer_id = int(layer_number) - 1 if layer_number is not None else int(match.group(1)) + if layer_id not in self.selected_layers: + continue + module.register_forward_hook(partial(self.record_layer, layer_id)) + registered_layers += 1 + for module_name, module in model_chunk.named_modules(): + if any(module_name.endswith(suffix) for suffix in self.module_suffixes): + module.register_forward_hook(partial(self.record_module, module_name)) + return registered_layers + + +def enable_megatron_layerwise_dump(args, model, store_prefix: str) -> None: + """Register one dump hook per selected layer on every model rank.""" + + dump_dir = os.getenv("VIME_LAYERWISE_ALIGNMENT_DUMP_DIR") + if not dump_dir: + return + + selected_layers = set(getattr(args, "megatron_deepgemm_forward_layers", []) or []) + if not selected_layers: + raise RuntimeError("VIME_LAYERWISE_ALIGNMENT_DUMP_DIR requires " "--megatron-deepgemm-forward-layers") + + registered_layers = 0 + for model_chunk in model: + if getattr(model_chunk, "_vime_layerwise_dump_registered", False): + continue + dumper = _MegatronLayerwiseDumper(dump_dir, selected_layers, store_prefix) + registered_layers += dumper.register(model_chunk) + model_chunk._vime_layerwise_dump_registered = True + model_chunk._vime_layerwise_dumper = dumper + + if registered_layers == 0 and not any( + getattr(model_chunk, "_vime_layerwise_dump_registered", False) for model_chunk in model + ): + raise RuntimeError("Could not find selected Megatron decoder layers to dump") + if registered_layers: + logger.info( + "Enabled Megatron layerwise alignment dump for layers %s at %s", + sorted(selected_layers), + dump_dir, + ) diff --git a/vime/backends/megatron_utils/arguments.py b/vime/backends/megatron_utils/arguments.py index 59a05dca7..5176eb8fc 100644 --- a/vime/backends/megatron_utils/arguments.py +++ b/vime/backends/megatron_utils/arguments.py @@ -3,9 +3,13 @@ from megatron.training.arguments import parse_args as _megatron_parse_args from megatron.training.arguments import validate_args as _megatron_validate_args -from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding from transformers import AutoConfig +try: + from megatron.core.tokenizers.utils.build_tokenizer import vocab_size_with_padding as _vocab_size_with_padding +except ImportError: + from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding + __all__ = ["validate_args", "megatron_parse_args", "set_default_megatron_args"] logger = logging.getLogger(__name__) @@ -147,12 +151,22 @@ def equal(x, y): def _set_default_megatron_args(args): # always use zero optimizer args.use_distributed_optimizer = True + if not hasattr(args, "enable_gloo_process_groups"): + args.enable_gloo_process_groups = True # TODO: maybe change this after megatron has good fp8 support args.bf16 = not args.fp16 + # Checkpoint I/O defaults: these keep checkpoint contents unchanged while + # reducing repeated validation/planning work and parallelizing load. + args.use_persistent_ckpt_worker = True + args.ckpt_assume_constant_structure = True + args.ckpt_fully_parallel_load = True # placeholders if args.seq_length is None: args.seq_length = 4096 - args.max_position_embeddings = args.seq_length + # Megatron also uses this value as YaRN's original context length. Preserve + # the checkpoint/model value when the launcher supplied one explicitly. + if args.max_position_embeddings is None: + args.max_position_embeddings = args.seq_length # TODO: revisit this when megatron(dev) have solved the optimizer-cpu-offload ckpt saving bug args.dist_ckpt_save_pre_mcore_014 = True # compatible for megatron diff --git a/vime/backends/megatron_utils/checkpoint.py b/vime/backends/megatron_utils/checkpoint.py index 1e6f9193a..5ee3c7de1 100644 --- a/vime/backends/megatron_utils/checkpoint.py +++ b/vime/backends/megatron_utils/checkpoint.py @@ -8,11 +8,7 @@ from megatron.training.checkpointing import save_checkpoint from megatron.training.global_vars import get_args -from vime.utils import megatron_bridge_utils -from vime.utils.common import is_npu - -logger = logging.getLogger(__name__) - +from vime.platforms import current_platform try: # Here we patch out the `validate_non_overlapping_shards_metadata` in both functions @@ -91,22 +87,17 @@ def _init_from_local_shards_and_global_metadata( # type: ignore[override] ShardedTensor._init_from_local_shards_and_global_metadata = _init_from_local_shards_and_global_metadata - if is_npu() and hasattr(default_planner, "_validate_global_plan"): - - def patched_validate_global_plan(global_plan, metadata): - logger.info("[Patch] Skipping validate_access_integrity") - return True - - default_planner._validate_global_plan = patched_validate_global_plan + current_platform().checkpoint.patch_default_planner(default_planner) except ImportError: pass +logger = logging.getLogger(__name__) __all__ = ["save_checkpoint"] -def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_context, skip_load_to_model_and_opt): +def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_context): # ref: how megatron `load_checkpoint` gets directory args = get_args() load_path = args.load @@ -121,7 +112,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_con optimizer=optimizer, opt_param_scheduler=opt_param_scheduler, checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=skip_load_to_model_and_opt, + skip_load_to_model_and_opt=False, ) else: return _load_checkpoint_hf( @@ -139,18 +130,10 @@ def _is_megatron_checkpoint(path: str | Path) -> bool: def _load_checkpoint_hf(ddp_model, optimizer, args, load_path: str): - assert args.megatron_to_hf_mode == "bridge", "Only bridge mode is supported for loading HF checkpoint" - from megatron.bridge import AutoBridge - - import vime_plugins.megatron_bridge # noqa: F401 - logger.info(f"Load checkpoint from HuggingFace model into Megatron (path={load_path})") + from vime.backends.megatron_utils.hf_to_megatron import load_hf_weights - with megatron_bridge_utils.patch_megatron_model(ddp_model): - bridge = megatron_bridge_utils.patch_auto_bridge_hf_config( - AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True) - ) - bridge.load_hf_weights(ddp_model) + load_hf_weights(args, ddp_model, load_path) # Copied from Megatron-core :: load_checkpoint (with simplifications) if (args.fp16 or args.bf16) and optimizer is not None: diff --git a/vime/backends/megatron_utils/cp_utils.py b/vime/backends/megatron_utils/cp_utils.py index 448c154c6..a54da3a7e 100644 --- a/vime/backends/megatron_utils/cp_utils.py +++ b/vime/backends/megatron_utils/cp_utils.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Sequence import torch import torch.distributed as dist @@ -9,8 +9,6 @@ def get_logits_and_tokens_offset_with_cp( total_length: int, response_length: int, - qkv_format: str = "thd", - max_seq_len: int | None = None, ): """ All offsets start from the begining of the prompt. @@ -20,11 +18,7 @@ def get_logits_and_tokens_offset_with_cp( assert cp_size > 1 prompt_length = total_length - response_length - if qkv_format == "thd": - chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) - else: - assert max_seq_len is not None, "max_seq_len must be provided for qkv_format=bshd" - chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) + chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) # the offset of 2 chunks chunk_0 = (cp_rank * chunk_size, (cp_rank + 1) * chunk_size) @@ -56,8 +50,6 @@ def get_sum_of_sample_mean( loss_masks: list[torch.Tensor], sample_denoms: list[torch.Tensor] | torch.Tensor | None = None, calculate_per_token_loss: bool = False, - qkv_format: str = "thd", - max_seq_lens: list[int] | None = None, ) -> Callable[[torch.Tensor], torch.Tensor]: """ Calculate correct sample mean for CP. @@ -100,18 +92,14 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: cp_chunk_lengths: list[int] = [] chunked_loss_masks: list[torch.Tensor] = [] - for i, (total_length, response_length, loss_mask) in enumerate( - zip(total_lengths, response_lengths, loss_masks, strict=False) - ): - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + for total_length, response_length, loss_mask in zip(total_lengths, response_lengths, loss_masks, strict=False): prompt_length = total_length - response_length - _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length, qkv_format, max_seq_len - ) + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length) loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] - chunked_loss_masks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0)) - cp_chunk_lengths.append(chunked_loss_masks[i].size(0)) + chunked_loss_mask = torch.cat([loss_mask_0, loss_mask_1], dim=0) + chunked_loss_masks.append(chunked_loss_mask) + cp_chunk_lengths.append(chunked_loss_mask.size(0)) def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( @@ -136,114 +124,6 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token -def reduce_train_step_metrics( - losses_reduced: list[dict], - *, - calculate_per_token_loss: bool, - step_global_batch_size: int, - cp_size: int, - dp_with_cp_group, -) -> dict[str, float]: - """Aggregate per-mb log dicts into the dict ``train_one_step`` reports. - - Pipeline (1:1 with what the train loop used to do inline): - 1. Sum each metric's per-mb ``values`` tensor locally on this rank. - 2. All-reduce across the DP*CP group (``dp_with_cp_group``). - 3. Apply the per-mode divisor / cp_factor: - - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``, - CP-inflated by ``cp_size`` because every CP rank computes the same - num_tokens off the FULL (not chunked) masks; the - ``cp_factor = cp_size`` multiplier cancels that inflation, leaving - the genuine per-token average. - - per-rollout-mean: divisor = constant ``step_global_batch_size`` from - the rollout side, never all-reduced, so no CP inflation to cancel - and ``cp_factor = 1``. - - Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce`` - to a no-op, then pre-aggregate virtual ranks themselves — this exercises - the same call shape as production while staying single-process. - """ - keys = losses_reduced[0]["keys"] - values = None - for x in losses_reduced: - values = x["values"] if values is None else values + x["values"] - assert len(keys) + 1 == values.numel() - dist.all_reduce(values, group=dp_with_cp_group) - values = values.tolist() - - if calculate_per_token_loss: - num_samples_or_tokens = values[0] - cp_factor = cp_size - else: - num_samples_or_tokens = step_global_batch_size - cp_factor = 1 - return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)} - - -def rollout_log_metric_contribution( - per_rank_reducer_sum: float, - *, - cp_size: int, - num_rollouts_in_rollout: int, - dp_size: int, -) -> tuple[float, float]: - """``(sum, count)`` tuple to hand the gather step for a per-rollout-mean - metric on the rollout side (``log_rollout_data``). - - Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout`` - (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size * - cp_size`` ranks, and each rank emits the same ``count``, so the totals - cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount = - sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports - for the same samples (when ``num_steps_per_rollout == 1``). - - Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end - in tests (single helper call per rank, returns the reduced number on - the source rank). - """ - sum_value = cp_size * per_rank_reducer_sum - count = num_rollouts_in_rollout / dp_size - return sum_value, count - - -def gather_and_reduce_log_dict( - log_dict: dict, - *, - dp_size: int, - dp_src_rank: int, - dp_group, -) -> dict | None: - """``dist.gather_object`` per-rank log_dicts + per-key reduction. - - Per key in the gathered dicts: - - ``(sum, count)`` tuple → ``Σsum / Σcount`` (per-rollout-mean shape; - pair with :func:`rollout_log_metric_contribution`). - - plain value → ``Σ / dp_size`` (legacy mean-across-ranks; the only - correct answer when ranks hold the same data). - - Returns the reduced dict on ``dp_src_rank``, ``None`` elsewhere. The - caller adds whatever metric-name prefix / wandb plumbing it wants — - this helper stays free of side effects so CPU multi-process unit tests - can drive it directly with real ``torch.distributed``. - """ - if dist.get_rank() == dp_src_rank: - gathered = [None] * dp_size - dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group) - reduced: dict = {} - for key in log_dict: - values = [d[key] for d in gathered] - first = values[0] - if isinstance(first, tuple) and len(first) == 2: - total_sum = sum(v[0] for v in values) - total_count = sum(v[1] for v in values) - reduced[key] = total_sum / total_count if total_count else 0.0 - else: - reduced[key] = sum(values) / dp_size - return reduced - dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group) - return None - - def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor: """ Gather tensors across all ranks in the context parallel group. @@ -299,15 +179,10 @@ def zero(len: int) -> torch.Tensor: def slice_with_cp( tokens: torch.Tensor, pad_value: tuple[int, float, Callable], - qkv_format: str = "thd", - max_seq_len: int | None = None, ) -> torch.Tensor: cp_rank = mpu.get_context_parallel_rank() cp_size = mpu.get_context_parallel_world_size() - if qkv_format == "bshd": - assert max_seq_len is not None - def pad_tokens(tokens, pad): if isinstance(pad_value, Callable): pad_func = pad_value @@ -319,16 +194,10 @@ def pad_tokens(tokens, pad): return tokens if cp_size == 1: - if qkv_format == "bshd": - pad = max_seq_len - tokens.size(0) - tokens = pad_tokens(tokens, pad) return tokens token_len = len(tokens) - if qkv_format == "thd": - chunk_size = (token_len + 2 * cp_size - 1) // (2 * cp_size) - else: - chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) + chunk_size = (token_len + 2 * cp_size - 1) // (2 * cp_size) # pad pad = 2 * cp_size * chunk_size - token_len @@ -344,8 +213,6 @@ def slice_log_prob_with_cp( log_prob: list[float] | torch.Tensor, total_length: int, response_length: int, - qkv_format: str = "thd", - max_token_len: int | None = None, ) -> list[float] | torch.Tensor: assert len(log_prob) == response_length, ( f"log_prob length mismatch: len(log_prob)={len(log_prob)}, " @@ -358,9 +225,7 @@ def slice_log_prob_with_cp( return log_prob prompt_length = total_length - response_length - _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp( - total_length, response_length, qkv_format, max_token_len - ) + _, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length) chunk_1 = log_prob[logits_offset[0][0] - (prompt_length - 1) : logits_offset[0][1] - (prompt_length - 1)] chunk_2 = log_prob[logits_offset[1][0] - (prompt_length - 1) : logits_offset[1][1] - (prompt_length - 1)] @@ -369,3 +234,64 @@ def slice_log_prob_with_cp( return chunk_1 + chunk_2 else: return torch.cat([chunk_1, chunk_2], dim=0) + + +def _pad_routed_experts(experts: torch.Tensor, pad: int, num_experts: int) -> torch.Tensor: + if pad == 0: + return experts + _, num_layers, topk = experts.shape + pad_experts = ( + torch.arange( + pad * num_layers * topk, + device=experts.device, + dtype=experts.dtype, + ).reshape((pad, num_layers, topk)) + % num_experts + ) + return torch.cat([experts, pad_experts], dim=0) + + +def prepare_routed_experts_for_routing_replay( + rollout_routed_experts: Sequence[torch.Tensor], + tokens: Sequence[torch.Tensor], + *, + num_experts: int, + data_pad_size_multiplier: int, + sequence_parallel: bool, + allgather_cp: bool, +) -> torch.Tensor: + """Align rollout routed-experts metadata with the training token layout.""" + assert len(rollout_routed_experts) == len(tokens) + for experts, token_ids in zip(rollout_routed_experts, tokens, strict=False): + assert experts.shape[0] == token_ids.shape[0] - 1, f"{experts.shape}, {token_ids.shape}" + + padded_experts = [_pad_routed_experts(experts, 1, num_experts) for experts in rollout_routed_experts] + pad_size = mpu.get_tensor_model_parallel_world_size() * data_pad_size_multiplier + + if allgather_cp: + routed_experts = torch.cat(padded_experts, dim=0) + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + global_pad_size = cp_size * pad_size + pad = (global_pad_size - routed_experts.size(0) % global_pad_size) % global_pad_size + routed_experts = _pad_routed_experts(routed_experts, pad, num_experts) + routed_experts = routed_experts.chunk(cp_size, dim=0)[cp_rank] + else: + routed_experts = [ + slice_with_cp(experts, lambda x, pad: _pad_routed_experts(x, pad, num_experts)) + for experts in padded_experts + ] + routed_experts = torch.cat(routed_experts, dim=0) + pad = (pad_size - routed_experts.size(0) % pad_size) % pad_size + routed_experts = _pad_routed_experts(routed_experts, pad, num_experts) + + if sequence_parallel: + tp_rank = mpu.get_tensor_model_parallel_rank() + tp_size = mpu.get_tensor_model_parallel_world_size() + seqlen = routed_experts.size(0) + assert seqlen % tp_size == 0 + start = seqlen // tp_size * tp_rank + end = seqlen // tp_size * (tp_rank + 1) + routed_experts = routed_experts[start:end] + + return routed_experts diff --git a/vime/backends/megatron_utils/data.py b/vime/backends/megatron_utils/data.py index 42c19e7e6..7cd1bce5d 100644 --- a/vime/backends/megatron_utils/data.py +++ b/vime/backends/megatron_utils/data.py @@ -1,35 +1,20 @@ -import logging -from argparse import Namespace from collections.abc import Sequence -import numpy as np import torch -import torch.distributed as dist import torch.nn.functional as F from megatron.core import mpu from megatron.core.packed_seq_params import PackedSeqParams -from vime.utils import train_metric_utils -from vime.utils.flops_utils import calculate_fwd_flops -from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step +from vime.utils import accelerator from vime.utils.types import RolloutBatch -from ...utils import logging_utils -from .cp_utils import ( - gather_and_reduce_log_dict, - get_sum_of_sample_mean, - rollout_log_metric_contribution, - slice_with_cp, -) - -logger = logging.getLogger(__name__) +from .cp_utils import slice_with_cp def get_batch( data_iterator: "DataIterator", keys: Sequence[str], pad_multiplier: int = 128, - qkv_format: str = "thd", allgather_cp: bool = False, ) -> dict[str, torch.Tensor | PackedSeqParams | list[torch.Tensor] | None]: """ @@ -67,63 +52,53 @@ def get_batch( cp_size = mpu.get_context_parallel_world_size() cp_rank = mpu.get_context_parallel_rank() - if qkv_format == "bshd": - max_seqlen = batch["max_seq_lens"][0] - assert max([t.size(0) for t in tokens]) <= max_seqlen - tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] - tokens = torch.stack(tokens) - packed_seq_params = None + if allgather_cp: + # DSA mode: concatenate all sequences first, then slice once with CP. + # We also pad the *global* concatenated stream to make per-rank chunks equal. + cu_seqlens_list: list[int] = [0] + for t in tokens: + cu_seqlens_list.append(cu_seqlens_list[-1] + t.size(0)) - elif qkv_format == "thd": - if allgather_cp: - # DSA mode: concatenate all sequences first, then slice once with CP. - # We also pad the *global* concatenated stream to make per-rank chunks equal. - cu_seqlens_list: list[int] = [0] - for t in tokens: - cu_seqlens_list.append(cu_seqlens_list[-1] + t.size(0)) - - tokens = torch.cat(tokens, dim=0) - - # Pad global stream so (1) divisible by cp_size (equal chunks), - # (2) divisible by pad_size (reduce fragmentation). - global_pad_size = cp_size * pad_size - pad = (global_pad_size - tokens.size(0) % global_pad_size) % global_pad_size - if pad != 0: - tokens = F.pad(tokens, (0, pad), value=pad_token_id) - cu_seqlens_list.append(cu_seqlens_list[-1] + pad) - - cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device()) - tokens = tokens.chunk(cp_size, dim=0)[cp_rank] - else: - tokens = [slice_with_cp(t, pad_token_id, qkv_format) for t in tokens] - - cu_seqlens = [0] - for t in tokens: - cu_seqlens.append(cu_seqlens[-1] + t.size(0)) - - tokens = torch.cat(tokens) - - # Always pad to reduce memory fragmentation and maybe make the computation faster - pad = (pad_size - tokens.size(0) % pad_size) % pad_size - if pad != 0: - tokens = F.pad(tokens, (0, pad), value=pad_token_id) - cu_seqlens.append(cu_seqlens[-1] + pad) - - # thd requires the cu_seqlens to be of the origin length - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size - - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - packed_seq_params = PackedSeqParams( - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, - qkv_format="thd", - ) - - tokens = tokens.unsqueeze(0) + tokens = torch.cat(tokens, dim=0) + + # Pad global stream so (1) divisible by cp_size (equal chunks), + # (2) divisible by pad_size (reduce fragmentation). + global_pad_size = cp_size * pad_size + pad = (global_pad_size - tokens.size(0) % global_pad_size) % global_pad_size + if pad != 0: + tokens = F.pad(tokens, (0, pad), value=pad_token_id) + cu_seqlens_list.append(cu_seqlens_list[-1] + pad) + + cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=accelerator.current_device()) + tokens = tokens.chunk(cp_size, dim=0)[cp_rank] else: - raise ValueError(f"Unsupported qkv_format: {qkv_format}") + tokens = [slice_with_cp(t, pad_token_id) for t in tokens] + + cu_seqlens = [0] + for t in tokens: + cu_seqlens.append(cu_seqlens[-1] + t.size(0)) + + tokens = torch.cat(tokens) + + # Always pad to reduce memory fragmentation and maybe make the computation faster + pad = (pad_size - tokens.size(0) % pad_size) % pad_size + if pad != 0: + tokens = F.pad(tokens, (0, pad), value=pad_token_id) + cu_seqlens.append(cu_seqlens[-1] + pad) + + # thd requires the cu_seqlens to be of the origin length + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int, device=accelerator.device()) * cp_size + + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + qkv_format="thd", + ) + + tokens = tokens.unsqueeze(0) batch["tokens"] = tokens batch["packed_seq_params"] = packed_seq_params @@ -142,75 +117,41 @@ def get_batch( if allgather_cp: loss_masks.append(loss_mask) continue - loss_mask = slice_with_cp(loss_mask, 0, qkv_format, max_seqlen) + loss_mask = slice_with_cp(loss_mask, 0) loss_masks.append(loss_mask) - if qkv_format == "bshd": - loss_masks = torch.stack(loss_masks) - elif qkv_format == "thd" and allgather_cp: + if allgather_cp: # DSA: concatenate first (same as tokens), pad globally (same pad as above), then slice once. loss_masks = torch.cat(loss_masks, dim=0) if pad != 0: loss_masks = F.pad(loss_masks, (0, pad), value=0) loss_masks = loss_masks.chunk(cp_size, dim=0)[cp_rank].unsqueeze(0) - elif qkv_format == "thd": + else: loss_masks = torch.cat(loss_masks) loss_masks = F.pad(loss_masks, (0, pad), value=0).unsqueeze(0) assert loss_masks.shape == tokens.shape, f"loss_masks.shape: {loss_masks.shape}, tokens.shape: {tokens.shape}" batch["full_loss_masks"] = loss_masks - # Process multimodal training tensors if present multimodal_train_inputs = batch.get("multimodal_train_inputs", None) if multimodal_train_inputs is not None: - multimodal_data = {} # key -> concatenated tensor + multimodal_data = {} for mm_input_dict in multimodal_train_inputs: if mm_input_dict is not None: for key, mm_tensor in mm_input_dict.items(): - if key not in multimodal_data: - multimodal_data[key] = mm_tensor - else: - multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0) + mm_tensor = torch.atleast_1d(torch.as_tensor(mm_tensor)) + if key in multimodal_data: + current = multimodal_data[key] + max_len = max(current.shape[-1], mm_tensor.shape[-1]) + mm_tensor = torch.cat( + [F.pad(tensor, (0, max_len - tensor.shape[-1])) for tensor in (current, mm_tensor)] + ) + multimodal_data[key] = mm_tensor batch["multimodal_train_inputs"] = multimodal_data return batch -def gather_log_data( - metric_name: str, - args: Namespace, - rollout_id: int, - log_dict: dict[str, "float | tuple[float, float]"], -) -> dict[str, float] | None: - """ - Gather per-rank metrics, reduce on the DP source rank, and log to W&B / TB. - - Each value in ``log_dict`` is either: - * a ``(sum, count)`` tuple → reduced as ``Σsum / Σcount``; - * a plain scalar → reduced as ``Σ / dp_size`` (mean across ranks). - - The gather + reduce step is delegated to - :func:`cp_utils.gather_and_reduce_log_dict` so it can be exercised by - CPU multi-process unit tests directly. This function adds the - ``metric_name`` prefix and the W&B / TB logging side effects. - """ - reduced = gather_and_reduce_log_dict( - log_dict, - dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True), - dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True), - dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), - ) - if reduced is None: - return None - reduced_log_dict = {f"{metric_name}/{k}": v for k, v in reduced.items()} - logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") - # Calculate step once to avoid duplication - step = compute_rollout_step(args, rollout_id) - reduced_log_dict["rollout/step"] = step - logging_utils.log(args, reduced_log_dict, step_key="rollout/step") - return reduced_log_dict - - class DataIterator: """Iterator over a rollout dict following an explicit micro-batch index schedule.""" @@ -258,271 +199,6 @@ def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]: return [DataIterator(rollout_data, micro_batch_indices) for _ in range(vpp_size)] -def log_rollout_data( - rollout_id: int, - args: Namespace, - rollout_data: RolloutBatch, -) -> None: - """ - Summarize rollout fields and log reduced metrics on PP last stage, TP rank 0. - - - Tensor-valued lists are concatenated and averaged. For token-level metrics - like log-probs/returns/advantages/values, computes a CP-correct sample mean - using `loss_masks` and total/response lengths. - - Non-tensor lists are averaged elementwise. - - Scalars are converted to Python numbers. - """ - if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): - cp_size = mpu.get_context_parallel_world_size() - log_dict = {} - response_lengths = rollout_data["response_lengths"] - loss_masks = rollout_data["loss_masks"] - total_lengths = rollout_data["total_lengths"] - max_seq_lens = rollout_data.get("max_seq_lens", None) - # Same per-rollout denominators the training loss uses, so reported - # log_probs / returns / advantages / etc. live in the same per-rollout - # mean space (rather than per-sample) as the gradient signal. - rollout_mask_sums = rollout_data.get("rollout_mask_sums", None) - # For per-rollout-mean metrics: ``rollout_log_metric_contribution`` - # produces the ``(sum, count)`` tuple so gather_log_data's - # ``Σsum / Σcount`` lands on ``sum_DP_full / num_rollouts`` — the - # same number train_one_step reports for the same samples. - dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False) - num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"]) - - for key, val in rollout_data.items(): - if key in [ - "tokens", - "multimodal_train_inputs", - "loss_masks", - "sample_indices", - "rollout_ids", - "rollout_mask_sums", - "rollout_routed_experts", - "max_seq_lens", - "global_batch_sizes", - "num_microbatches", - "micro_batch_indices", - ]: - continue - # Emit (sum, count) so gather_log_data can do a weighted average across - # DP ranks. This stops the legacy "every rank has the same N samples" - # assumption from biasing means once uneven-DP partitioning lands. - if isinstance(val, (list, tuple)): - count = len(val) - if isinstance(val[0], torch.Tensor): - # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be - # modified in place and will cause problem for the next rollout. - if key in [ - "log_probs", - "ref_log_probs", - "rollout_log_probs", - "returns", - "advantages", - "values", - "teacher_log_probs", - "opd_reverse_kl", - ]: - tensor = torch.cat(val).clone().detach() - sum_of_sample_mean = get_sum_of_sample_mean( - total_lengths, - response_lengths, - loss_masks, - rollout_mask_sums, - qkv_format=args.qkv_format, - max_seq_lens=max_seq_lens, - ) - # Compute (sum, count) via the shared helper so this - # path and the unit tests stay in sync. - sum_value, count = rollout_log_metric_contribution( - sum_of_sample_mean(tensor).item(), - cp_size=cp_size, - num_rollouts_in_rollout=num_rollouts_in_rollout, - dp_size=dp_world, - ) - log_dict[key] = (sum_value, count) - continue - tensor = torch.cat(val).clone().detach() - # val.mean() * cp_size is the per-sample mean for one rank; - # multiply by count to get the per-rank sum. - per_rank_sum = tensor.mean() * cp_size * count - sum_value = per_rank_sum.item() - else: - sum_value = sum(val) - log_dict[key] = (sum_value, count) - elif isinstance(val, torch.Tensor): - # Scalar tensor (one per rank): treat as count=1. - log_dict[key] = (val.float().mean().item(), 1) - else: - raise ValueError(f"Unsupported type: {type(val)} for key: {key}") - - reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict) - if args.ci_test and reduced_log_dict is not None: - # This is an initial actor/ref zero-KL check. R3 replays rollout - # routing for the actor forward, while the reference forward - # intentionally falls through to normal routing, so their - # log-probs are not expected to match bit-for-bit in CI. - if ( - rollout_id == 0 - and not getattr(args, "ci_disable_kl_checker", False) - and not getattr(args, "use_rollout_routing_replay", False) - and "rollout/log_probs" in reduced_log_dict - and "rollout/ref_log_probs" in reduced_log_dict - ): - # TODO: figure out why there is a small numerical difference in log_probs and ref_log_probs in CI test, and whether it's expected or not. - # assert reduced_log_dict["rollout/log_probs"] == reduced_log_dict["rollout/ref_log_probs"] - assert abs(reduced_log_dict["rollout/log_probs"] - reduced_log_dict["rollout/ref_log_probs"]) < 1e-8 - if "rollout/log_probs" in reduced_log_dict: - assert -1 < reduced_log_dict["rollout/log_probs"] < 0 - if "rollout/entropy" in reduced_log_dict: - assert 0 < reduced_log_dict["rollout/entropy"] < 1 - - if args.log_multi_turn: - log_multi_turn_data(rollout_id, args, rollout_data) - if args.log_passrate: - log_passrate(rollout_id, args, rollout_data) - - if args.log_correct_samples: - if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): - cp_size = mpu.get_context_parallel_world_size() - log_dict = {} - response_lengths = rollout_data["response_lengths"] - loss_masks = rollout_data["loss_masks"] - total_lengths = rollout_data["total_lengths"] - - def quantile(total_value, n_quantiles, data) -> dict: - import math - - assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1." - - quantiles = [((i + 1) / n_quantiles) for i in range(n_quantiles)] - cut_points = [total_value * q for q in quantiles] - cut_points[-1] = total_value - - count = [0] * n_quantiles - for d in data: - for i, point in enumerate(cut_points): - if d <= point: - count[i] += 1 - break - - total = sum(count) + 1e-9 - percentile = [c / total for c in count] - - percentile = {f"p{min(math.ceil(q*100),100)}": p for q, p in zip(quantiles, percentile, strict=True)} - return percentile - - raw_rewards = rollout_data["raw_reward"] - # Additional metrics for correct cases are calculated separately below. - correct_response_lengths = [] - correct_total_lengths = [] - correct_loss_masks = [] - correct_entropy = [] - for i, raw_reward in enumerate(raw_rewards): - if raw_reward == 1: - correct_response_lengths.append(response_lengths[i]) - correct_total_lengths.append(total_lengths[i]) - correct_loss_masks.append(loss_masks[i]) - correct_entropy.append(-rollout_data["log_probs"][i]) - num_correct_responses = len(correct_total_lengths) - rollout_data["correct_response_lengths"] = correct_response_lengths - correct_response_length_percentile = quantile( - args.rollout_max_response_len, 4, rollout_data["correct_response_lengths"] - ) - for p, val in correct_response_length_percentile.items(): - rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses - if len(correct_entropy) > 0: - # NOTE: per-sample-mean over the correct subset, not per-rollout. - # A rollout's siblings may not all be correct, and slicing - # ``rollout_mask_sums`` here would leave a denom that still - # includes incorrect siblings — meaningless for a "correct-only" - # entropy report. Per-sample-mean over the filtered subset is - # the cleanest semantic. - sum_of_sample_mean = get_sum_of_sample_mean( - correct_total_lengths, correct_response_lengths, correct_loss_masks, sample_denoms=None - ) - correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0)) - rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses - else: - rollout_data["correct_entropy"] = [0] * num_correct_responses - - -def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None: - """ - Log multi-turn auxiliary metrics such as raw/observed response lengths and rounds. - - Operates only on PP last stage and TP rank 0. Uses GPU tensors when available - to compute statistics without host transfers. - """ - if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): - log_dict = {} - for key, val in rollout_data.items(): - if key == "loss_masks": - if val: # Check if val is not empty - device = val[0].device # Get device from first tensor - - # Vectorized length calculation using torch - raw_response_lengths = torch.tensor([v.shape[0] for v in val], dtype=torch.float32, device=device) - log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item() - log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item() - log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item() - log_dict["raw_response_length/response_length_clip_ratio"] = ( - (raw_response_lengths >= args.rollout_max_response_len).float().mean().item() - ) - - # Vectorized sum calculation using torch - stay on GPU - wo_obs_response_lengths = torch.tensor( - [v.sum().item() for v in val], dtype=torch.float32, device=device - ) - log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item() - log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item() - log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item() - if key == "round_number": - # Use numpy for vectorized round number statistics - round_number_array = np.array(val) - log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array) - log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array) - log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array) - gather_log_data("multi_turn", args, rollout_id, log_dict) - - -def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None: - """ - Compute pass@k metrics from `raw_reward` groups and log the results. - - `raw_reward` is reshaped to `[group_number, group_size]`, then pass@k is - estimated per problem and averaged. - """ - if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): - log_dict = {} - for key, val in rollout_data.items(): - if key != "raw_reward": - continue - - log_dict |= compute_pass_rate( - flat_rewards=val, - group_size=args.n_samples_per_prompt, - num_groups=args.rollout_batch_size, - ) - - gather_log_data("passrate", args, rollout_id, log_dict) - - -def log_perf_data(rollout_id: int, args: Namespace) -> None: - train_metric_utils.log_perf_data_raw( - rollout_id=rollout_id, - args=args, - is_primary_rank=( - mpu.get_tensor_model_parallel_rank() == 0 - and mpu.is_pipeline_last_stage() - and mpu.get_data_parallel_rank(with_context_parallel=True) == 0 - ), - compute_total_fwd_flops=lambda seq_lens: calculate_fwd_flops(seqlens=seq_lens, args=args) - / dist.get_world_size() - / 1e12, - ) - - def tensors_to_cpu(tensor_list): """Move a list of GPU tensors to CPU for Ray object store transfer. @@ -550,5 +226,5 @@ def tensors_to_gpu(tensor_list, device=None): if tensor_list is None: return None if device is None: - device = torch.cuda.current_device() + device = accelerator.current_device() return [t.to(device=device, dtype=torch.float32) for t in tensor_list] diff --git a/vime/backends/megatron_utils/dspark/__init__.py b/vime/backends/megatron_utils/dspark/__init__.py new file mode 100644 index 000000000..ef091be9f --- /dev/null +++ b/vime/backends/megatron_utils/dspark/__init__.py @@ -0,0 +1 @@ +"""Megatron adaptation of DeepSpec's DSpark draft model.""" diff --git a/vime/backends/megatron_utils/dspark/attention.py b/vime/backends/megatron_utils/dspark/attention.py new file mode 100644 index 000000000..2378260b4 --- /dev/null +++ b/vime/backends/megatron_utils/dspark/attention.py @@ -0,0 +1,203 @@ +"""DSpark dual-input attention for Megatron backend. + +Adapted from DeepSpec/deepspec/modeling/dspark/qwen3/modeling.py:Qwen3DSparkAttention. + +Key difference from standard attention: K and V are computed from BOTH the +draft hidden states AND the target (policy) hidden states, then concatenated: + + k = cat([k_proj(target_hidden), k_proj(draft_hidden)], dim=seq) + v = cat([v_proj(target_hidden), v_proj(draft_hidden)], dim=seq) + +This dual-input K/V is the core architectural feature of DSpark that enables +the draft model to attend to the policy's intermediate representations. + +The draft model is replicated on every TP rank and uses plain ``nn.Linear`` +plus SDPA because Megatron TE attention does not support dual-input K/V. +""" + +import torch +import torch.nn.functional as F +from torch import nn + + +def apply_rotary_pos_emb(q, k, cos, sin): + """Apply rotary embeddings to q and k. + + Args: + q: [bsz, heads, q_len, head_dim] + k: [bsz, kv_heads, kv_len, head_dim] + cos: [bsz, 1, kv_len, head_dim] (covers full kv sequence) + sin: [bsz, 1, kv_len, head_dim] + Returns: + q_embed, k_embed (same shapes as q, k) + """ + q_len = q.size(-2) + # Q only attends to the last q_len positions of the rotary table + # (draft positions are after context positions) + q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def rotate_half(x): + """Rotate the second half of the last dim to the front (inverse concat).""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +class DSparkRotaryEmbedding(nn.Module): + """Precompute rotary sin/cos for DSpark positions. + + DSpark position ids cover both context (0..seq_len-1) and draft tokens + (anchor_pos..anchor_pos+block_size-1 per block). The rotary table must + cover the max position id, which is max(anchor_pos) + block_size - 1. + """ + + def __init__(self, head_dim: int, rotary_base: float = 10000.0): + super().__init__() + self.head_dim = head_dim + self.rotary_base = rotary_base + # Precompute a large table; will be indexed as needed. + # Max position ~ seq_len + num_anchors * block_size, which is bounded + # by the model's max_position_embeddings. + inv_freq = 1.0 / (rotary_base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute (cos, sin) for the given position ids. + + Args: + position_ids: [bsz, seq_len] long tensor + Returns: + cos, sin: [bsz, 1, seq_len, head_dim] each (broadcast-ready) + """ + # inv_freq: [head_dim/2] + # position_ids: [bsz, seq_len] + # freqs: [bsz, seq_len, head_dim/2] + inv_freq = self.inv_freq.float() # [head_dim/2] + positions = position_ids.float() # [bsz, seq_len] + # Outer product per batch: [bsz, seq_len, head_dim/2] + freqs = torch.einsum("i,bj->bji", inv_freq, positions) + emb = torch.cat([freqs, freqs], dim=-1) # [bsz, seq_len, head_dim] + cos = emb.cos() + sin = emb.sin() + + # Add head dim for broadcasting: [bsz, 1, seq_len, head_dim] + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return cos.to(position_ids.dtype), sin.to(position_ids.dtype) + + +class DSparkParallelAttention(nn.Module): + """Dual-input attention for DSpark draft model. + + K and V are computed from both ``hidden_states`` (draft) and + ``target_hidden_states`` (policy), then concatenated along the sequence + dimension. Q is computed from ``hidden_states`` only. + + The projections are intentionally unsharded because the complete draft + model runs on every TP rank. + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + attention_bias: bool = False, + rms_norm_eps: float = 1e-6, + rotary_base: float = 10000.0, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.num_key_value_groups = num_attention_heads // num_key_value_heads + self.head_dim = head_dim + self.scaling = head_dim**-0.5 + self.attention_bias = attention_bias + + self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=attention_bias) + self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias) + self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias) + self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=attention_bias) + + # QK-norm (Qwen3-style RMSNorm on per-head Q/K) + self.q_norm = nn.RMSNorm(head_dim, eps=rms_norm_eps) + self.k_norm = nn.RMSNorm(head_dim, eps=rms_norm_eps) + + self.rotary_emb = DSparkRotaryEmbedding(head_dim, rotary_base=rotary_base) + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass. + + Args: + hidden_states: [bsz, q_len, hidden] draft hidden states + target_hidden_states: [bsz, ctx_len, hidden] policy hidden states + position_ids: [bsz, ctx_len + q_len] position ids for rotary + attention_mask: [bsz, 1, q_len, ctx_len + q_len] bool/float mask + (True = attend, False = masked). If float, -inf for masked. + Returns: + [bsz, q_len, hidden] attention output + """ + bsz, q_len, _ = hidden_states.shape + ctx_len = target_hidden_states.shape[1] + + # Q from draft hidden states + q = self.q_proj(hidden_states).view(bsz, q_len, self.num_attention_heads, self.head_dim) + q = self.q_norm(q).transpose(1, 2) # [bsz, heads, q_len, head_dim] + + # K/V from BOTH target and draft, concatenated + k_ctx = self.k_proj(target_hidden_states) # [bsz, ctx_len, kv_heads * head_dim] + k_noise = self.k_proj(hidden_states) # [bsz, q_len, kv_heads * head_dim] + v_ctx = self.v_proj(target_hidden_states) + v_noise = self.v_proj(hidden_states) + + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, self.num_key_value_heads, self.head_dim) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, self.num_key_value_heads, self.head_dim) + k = self.k_norm(k).transpose(1, 2) # [bsz, kv_heads, kv_len, head_dim] + v = v.transpose(1, 2) # [bsz, kv_heads, kv_len, head_dim] + + # Apply rotary embeddings + # position_ids: [bsz, ctx_len + q_len] + cos, sin = self.rotary_emb(position_ids) + # cos/sin: [1, 1, seq_len, head_dim] but we need to match k's shape + # k shape: [bsz, kv_heads, kv_len, head_dim] + # cos shape: [1, 1, kv_len, head_dim] -> broadcast over bsz and heads + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + # Repeat K/V for GQA (grouped query attention) + if self.num_key_value_groups > 1: + k = k.repeat_interleave(self.num_key_value_groups, dim=1) + v = v.repeat_interleave(self.num_key_value_groups, dim=1) + + # SDPA attention + # attention_mask: [bsz, 1, q_len, kv_len] + # SDPA expects mask where True/1 = keep, False/0 = mask out (with bool mask) + # Or float mask where -inf = mask out + if attention_mask is not None: + if attention_mask.dtype == torch.bool: + # Convert to float bias: 0 for attend, -inf for mask + attn_bias = torch.zeros_like(attention_mask, dtype=q.dtype) + attn_bias = attn_bias.masked_fill(~attention_mask, float("-inf")) + else: + attn_bias = attention_mask + else: + attn_bias = None + + attn_output = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_bias, dropout_p=0.0, is_causal=False + ) # [bsz, heads, q_len, head_dim] + + attn_output = ( + attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.num_attention_heads * self.head_dim) + ) + return self.o_proj(attn_output) diff --git a/vime/backends/megatron_utils/dspark/common.py b/vime/backends/megatron_utils/dspark/common.py new file mode 100644 index 000000000..71481bb0b --- /dev/null +++ b/vime/backends/megatron_utils/dspark/common.py @@ -0,0 +1,330 @@ +"""DSpark common utilities: anchor sampling, mask construction, config. + +Adapted from DeepSpec/deepspec/modeling/dspark/common.py for vime Megatron backend. + +Key changes from DeepSpec: +- Removed ``add_metric`` calls (vime uses its own logging_utils). +- Replaced ``flex_attention.create_block_mask`` with an explicit SDPA-compatible + boolean attention mask because flex_attention is not available in all + Megatron/vLLM container images. +- Added ``DSparkConfig`` dataclass to bundle all DSpark hyperparameters. +""" + +from dataclasses import dataclass + +import torch +from torch import nn + + +@dataclass +class DSparkConfig: + """Bundle of DSpark hyperparameters passed to ``build_dspark_model``. + + Defaults match DeepSpec's ``config/dspark/dspark_qwen3_4b.py``. + """ + + # Backbone + block_size: int = 7 + num_draft_layers: int = 5 + target_layer_ids: tuple[int, ...] = (1, 9, 17, 25, 33) + mask_token_id: int = 151669 + num_anchors: int = 512 + + # Markov head + markov_rank: int = 256 + markov_head_type: str = "vanilla" + + # Confidence head + enable_confidence_head: bool = True + confidence_head_with_markov: bool = True + + # Loss weights + ce_loss_alpha: float = 0.1 + l1_loss_alpha: float = 0.9 + confidence_head_alpha: float = 1.0 + loss_decay_gamma: float = 4.0 + + # Draft loss combination weight (multiplies draft_loss added to policy_loss) + draft_loss_weight: float = 1.0 + + # Model dims (populated by build_dspark_model from policy config) + hidden_size: int = 0 + vocab_size: int = 0 + org_vocab_size: int = 0 # original (unpadded) vocab size for vLLM export + num_attention_heads: int = 0 + num_key_value_heads: int = 0 + head_dim: int = 0 + rms_norm_eps: float = 1e-6 + rotary_base: float = 10000.0 + rope_scaling: dict | None = None + + # MLP intermediate size (0 = auto-compute from hidden_size * 2.75) + # Set to match pre-trained checkpoint (e.g. 9728 for Qwen3-4B) + intermediate_size: int = 0 + + +@dataclass +class DSparkForwardOutput: + """Outputs for one DSpark training forward. + + Shapes: + draft_logits: [batch, num_anchors, block_size, vocab] + target_ids: [batch, num_anchors, block_size] + eval_mask: [batch, num_anchors, block_size] (bool) + block_keep_mask: [batch, num_anchors] (bool) + confidence_pred: [batch, num_anchors, block_size] (optional) + aligned_target_logits: [batch, num_anchors, block_size, vocab] (optional) + """ + + draft_logits: torch.Tensor + target_ids: torch.Tensor + eval_mask: torch.Tensor + block_keep_mask: torch.Tensor + confidence_pred: torch.Tensor | None = None + aligned_target_logits: torch.Tensor | None = None + + +class AcceptRatePredictor(nn.Module): + """Confidence head: predicts P(token accepted | prev accepted) per position.""" + + def __init__(self, input_dim: int): + super().__init__() + self.proj = nn.Linear(int(input_dim), 1) + + def forward(self, features): + return self.proj(features).squeeze(-1) + + +def validate_target_layer_ids(layer_ids, num_target_layers: int): + """Validate that target_layer_ids are strictly increasing and in range.""" + layer_ids = [int(layer_id) for layer_id in layer_ids] + assert layer_ids, "target_layer_ids must not be empty." + start = 0 + end = int(num_target_layers) - 1 + previous = None + for layer_id in layer_ids: + assert layer_id == -1 or start <= layer_id <= end, ( + f"target_layer_id {layer_id} is out of range {{-1}} U [{start}, {end}] " + f"for num_target_layers={num_target_layers}. " + "-1 denotes the embedding output." + ) + assert previous is None or layer_id > previous, "target_layer_ids must be strictly increasing." + previous = layer_id + return layer_ids + + +def build_anchor_candidate_mask( + *, + seq_len: int, + loss_mask: torch.Tensor, +) -> torch.Tensor: + """Build a boolean mask of valid anchor candidate positions. + + A position i is a valid anchor if both loss_mask[i] and loss_mask[i+1] are set + (the anchor token and its first prediction target must be supervised). + """ + num_candidates = max(seq_len - 1, 0) + if num_candidates == 0: + return loss_mask[:, :0].bool() + + anchor_valid = loss_mask[:, :num_candidates] > 0.5 + first_target_valid = loss_mask[:, 1 : num_candidates + 1] > 0.5 + return anchor_valid & first_target_valid + + +def sample_anchor_positions( + *, + seq_len: int, + loss_mask: torch.Tensor, + num_anchors: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample up to ``num_anchors`` anchor positions per sequence. + + Returns: + anchor_positions: [batch, num_anchors] long tensor (0 for invalid anchors) + block_keep_mask: [batch, num_anchors] bool tensor (True for valid anchors) + """ + valid = build_anchor_candidate_mask( + seq_len=seq_len, + loss_mask=loss_mask, + ) + valid_counts = valid.sum(dim=1) + bsz = loss_mask.shape[0] + num_candidates = valid.shape[1] + max_n = int(num_anchors) + if num_candidates == 0: + anchors = torch.zeros(bsz, max_n, dtype=torch.long, device=device) + keep_mask = torch.zeros(bsz, max_n, dtype=torch.bool, device=device) + return anchors, keep_mask + + indices = ( + torch.arange(num_candidates, device=device) + .unsqueeze(0) + .expand( + bsz, + -1, + ) + ) + masked_indices = torch.where( + valid, + indices, + torch.full_like(indices, seq_len + 1), + ) + random_vals = torch.rand(bsz, num_candidates, device=device) + random_vals = torch.where(valid, random_vals, torch.full_like(random_vals, 2.0)) + _, sorted_idx = random_vals.sort(dim=1) + gathered = torch.gather(masked_indices, 1, sorted_idx) + if num_candidates < max_n: + pad = torch.full( + (bsz, max_n - num_candidates), + seq_len + 1, + dtype=gathered.dtype, + device=device, + ) + gathered = torch.cat([gathered, pad], dim=1) + anchors = gathered[:, :max_n].sort(dim=1).values + keep_mask = torch.arange(max_n, device=device).unsqueeze(0) < (valid_counts.unsqueeze(1).clamp(max=max_n)) + anchors = torch.where(keep_mask, anchors, torch.zeros_like(anchors)) + return anchors, keep_mask + + +def build_eval_mask( + *, + seq_len: int, + loss_mask: torch.Tensor, + label_indices: torch.Tensor, + safe_label_indices: torch.Tensor, + block_keep_mask: torch.Tensor, +) -> torch.Tensor: + """Build the per-position evaluation mask. + + A position is evaluated if: + - its label index is within seq_len, + - its label position is enabled by loss_mask, + - its block is kept, + - and all preceding positions in the block are also evaluated (cumprod). + """ + target_valid = label_indices < seq_len + target_loss_mask = torch.gather( + loss_mask.unsqueeze(1).expand(-1, label_indices.size(1), -1), + 2, + safe_label_indices, + ) + eval_mask = target_valid & (target_loss_mask > 0.5) + eval_mask = eval_mask & block_keep_mask.unsqueeze(-1) + return eval_mask.to(torch.int32).cumprod(dim=-1).bool() + + +def create_position_ids( + anchor_positions: torch.Tensor, + block_size: int, +) -> torch.Tensor: + """Create position ids for draft tokens: [batch, num_blocks * block_size].""" + bsz, num_blocks = anchor_positions.shape + device = anchor_positions.device + offsets = torch.arange(block_size, device=device).view(1, 1, -1) + return (anchor_positions.unsqueeze(-1) + offsets).view( + bsz, + num_blocks * block_size, + ) + + +def create_noise_embed( + embed_tokens: nn.Module, + input_ids: torch.Tensor, + anchor_positions: torch.Tensor, + block_keep_mask: torch.Tensor, + *, + mask_token_id: int, + block_size: int, +) -> torch.Tensor: + """Create the noise embedding for DSpark draft tokens. + + Each block's first position is the anchor token; remaining positions are + the mask token. This is the input to the DSpark backbone. + """ + bsz = input_ids.shape[0] + num_blocks = anchor_positions.shape[1] + device = input_ids.device + noise_ids = torch.full( + (bsz, num_blocks * block_size), + mask_token_id, + dtype=torch.long, + device=device, + ) + block_starts = torch.arange(num_blocks, device=device) * block_size + block_starts = block_starts.unsqueeze(0).expand(bsz, -1) + anchor_tokens = torch.gather(input_ids, 1, anchor_positions) + flat_batch_idx = ( + torch.arange(bsz, device=device) + .unsqueeze(1) + .expand( + bsz, + num_blocks, + ) + ) + noise_ids[flat_batch_idx, block_starts] = torch.where( + block_keep_mask, + anchor_tokens, + torch.tensor(mask_token_id, dtype=torch.long, device=device), + ) + return embed_tokens(noise_ids) + + +def create_dspark_attention_mask( + *, + anchor_positions: torch.Tensor, + block_keep_mask: torch.Tensor, + seq_len: int, + block_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Build a dense boolean attention mask for DSpark. + + Replaces DeepSpec's flex_attention ``create_block_mask`` with an explicit + [batch, 1, q_len, kv_len] mask suitable for ``scaled_dot_product_attention``. + + Layout: + KV = [context (seq_len) | draft (num_blocks * block_size)] + Q = draft (num_blocks * block_size) + + For draft query in block b at position q_idx: + - attend to context positions [0, anchor_pos) + - attend to draft positions in the same block b + - masked out if block is not kept + """ + bsz, num_blocks = anchor_positions.shape + q_len = num_blocks * block_size + + # Query block id for each query position: [q_len] + q_block_id = torch.arange(q_len, device=device) // block_size + + # Anchor position per query position: [bsz, q_len] + anchor_per_q = anchor_positions[:, q_block_id] # [bsz, q_len] + + # Context mask: kv_idx < anchor_pos[bsz, q_idx] + # kv_ctx_idx: [seq_len], anchor_per_q: [bsz, q_len] + # Result: [bsz, q_len, seq_len] + kv_ctx_idx = torch.arange(seq_len, device=device) # [seq_len] + ctx_mask = kv_ctx_idx.unsqueeze(0).unsqueeze(0) < anchor_per_q.unsqueeze(-1) + + # Draft mask: kv_idx >= seq_len AND same block as query + kv_draft_idx = torch.arange(q_len, device=device) # [q_len] + kv_block_id = kv_draft_idx // block_size # [q_len] + # [q_len, q_len] + draft_mask = q_block_id.unsqueeze(1) == kv_block_id.unsqueeze(0) + # Expand to [bsz, q_len, q_len] + draft_mask = draft_mask.unsqueeze(0).expand(bsz, -1, -1) + + # block keep mask: [bsz, num_blocks] -> [bsz, q_len] + block_keep_per_q = block_keep_mask[:, q_block_id] # [bsz, q_len] + + # Combine: [bsz, q_len, kv_len] + full_mask = torch.cat([ctx_mask, draft_mask], dim=-1) # [bsz, q_len, kv_len] + full_mask = full_mask & block_keep_per_q.unsqueeze(-1) + + # Add head dim: [bsz, 1, q_len, kv_len] + full_mask = full_mask.unsqueeze(1) + return full_mask.to(dtype=dtype, device=device) diff --git a/vime/backends/megatron_utils/dspark/export.py b/vime/backends/megatron_utils/dspark/export.py new file mode 100644 index 000000000..3932e2c5b --- /dev/null +++ b/vime/backends/megatron_utils/dspark/export.py @@ -0,0 +1,65 @@ +"""Export DSpark weights with the names expected by vLLM.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +from torch import nn + +_SHARED_PARAM_NAMES = {"embed_tokens.weight", "lm_head.weight"} + + +def export_dspark_model_weights( + model_chunks: Sequence[nn.Module], + *, + use_policy_embedding: bool, +) -> list[tuple[str, torch.Tensor]]: + """Export draft weights and strip Megatron's vocabulary padding.""" + from megatron.core.utils import unwrap_model + + policy_model = None + for chunk in reversed(model_chunks): + model = unwrap_model(chunk) + if getattr(model, "draft_model", None) is not None: + policy_model = model + break + if policy_model is None: + raise RuntimeError("DSpark is enabled, but no draft model is attached to the policy") + + draft_model = policy_model.draft_model + config = getattr(draft_model, "config", None) + org_vocab_size = getattr(config, "org_vocab_size", 0) or 0 + padded_vocab_size = getattr(config, "vocab_size", 0) or 0 + + def strip_vocab_padding(tensor: torch.Tensor) -> torch.Tensor: + if padded_vocab_size > org_vocab_size > 0 and tensor.dim() >= 1 and tensor.shape[0] == padded_vocab_size: + return tensor[:org_vocab_size].contiguous() + return tensor + + hf_state = [ + (name, strip_vocab_padding(param)) + for name, param in draft_model.named_parameters() + if name not in _SHARED_PARAM_NAMES + ] + + embed = ( + _get_policy_embedding(policy_model) + if use_policy_embedding + else getattr(getattr(draft_model, "embed_tokens", None), "weight", None) + ) + if embed is not None: + hf_state.append(("embed_tokens.weight", strip_vocab_padding(embed))) + + lm_head = getattr(draft_model, "lm_head", None) + if lm_head is not None: + hf_state.append(("lm_head.weight", strip_vocab_padding(lm_head.weight))) + + return hf_state + + +def _get_policy_embedding(policy_model: nn.Module) -> torch.nn.Parameter | None: + embed = getattr(policy_model, "embedding", None) + if embed is None: + return None + return getattr(embed, "word_embeddings", embed).weight diff --git a/vime/backends/megatron_utils/dspark/hidden_capture.py b/vime/backends/megatron_utils/dspark/hidden_capture.py new file mode 100644 index 000000000..32d9bdb18 --- /dev/null +++ b/vime/backends/megatron_utils/dspark/hidden_capture.py @@ -0,0 +1,367 @@ +"""Hidden state capture for DSpark draft model training. + +Adapted from NeMo RL's ``nemo_rl/models/megatron/draft/hidden_capture.py``. + +During the policy forward pass, forward hooks capture: + 1. Embedding output (input_embeds) — from the policy's embedding layer + 2. Hidden states at ``target_layer_ids`` — concatenated for DSpark's fc projection + 3. Last layer hidden states — for L_tv / L_conf loss computation + +All captured states are gathered to the last pipeline stage (where the draft +model runs). PP=1 skips the send/recv path entirely. + +Key difference from NeMo RL Eagle3: + - Eagle3 captures ``hidden_states`` (concatenated aux layers) + ``inputs_embeds`` + - DSpark additionally captures ``target_last_hidden_states`` (policy's final layer) + because L_tv and L_conf require the policy's prediction at each draft position. +""" + +import logging +from contextlib import contextmanager +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor + +logger = logging.getLogger(__name__) + + +# Dtype encoding for send/recv (matches NeMo RL) +_DTYPE_TO_CODE = { + torch.float16: 0, + torch.bfloat16: 1, + torch.float32: 2, +} +_CODE_TO_DTYPE = {code: dtype for dtype, code in _DTYPE_TO_CODE.items()} + + +@dataclass +class CapturedStates: + """Container for hidden states captured from the policy model. + + Attributes: + target_hidden_states: [seq_len, bsz, num_target_layers * hidden] + Concatenated hidden states from policy's target_layer_ids. + inputs_embeds: [seq_len, bsz, hidden] embedding output (not used by + DSpark training forward, but captured for potential debugging). + target_last_hidden_states: [seq_len, bsz, hidden] policy's final layer + hidden states. Used for L_tv / L_conf computation. + """ + + target_hidden_states: Tensor | None = None + inputs_embeds: Tensor | None = None + target_last_hidden_states: Tensor | None = None + + +class HiddenStateCapture: + """Capture policy embeddings, aux-layer hidden states, and last-layer hidden states. + + This class registers forward hooks on the policy model's embedding layer, + target layers (specified by ``target_layer_ids``), and the last decoder + layer. After the policy forward pass, ``get_captured_states()`` returns + the gathered tensors. + + For PP>1, hidden states from earlier stages are sent to the last stage + via point-to-point send/recv. PP=1 skips this path. + """ + + def __init__( + self, + model: torch.nn.Module, + target_layer_ids: tuple[int, ...], + last_layer_idx: int | None = None, + ): + """Initialize the capture. + + Args: + model: the policy model (will be unwrapped if DDP-wrapped) + target_layer_ids: global layer indices to capture (0-indexed) + last_layer_idx: global index of the last layer to capture for + L_tv/L_conf. If None, uses ``max(target_layer_ids)`` or the + model's num_layers - 1. + """ + from megatron.core.utils import unwrap_model + + self.model = unwrap_model(model) + self.target_layer_ids = tuple(int(i) for i in target_layer_ids) + + # Determine number of layers + if hasattr(self.model, "decoder") and hasattr(self.model.decoder, "layers"): + self.num_layers = len(self.model.decoder.layers) + self._decoder = self.model.decoder + elif hasattr(self.model, "module"): + inner = self.model.module + if hasattr(inner, "decoder") and hasattr(inner.decoder, "layers"): + self.num_layers = len(inner.decoder.layers) + self._decoder = inner.decoder + self.model = inner + else: + raise RuntimeError("Cannot find decoder.layers in policy model. " f"Model type: {type(self.model)}") + else: + raise RuntimeError(f"Cannot find decoder.layers in policy model. Model type: {type(self.model)}") + + if last_layer_idx is None: + self.last_layer_idx = self.num_layers - 1 + else: + self.last_layer_idx = int(last_layer_idx) + + # PP info + try: + from megatron.core import parallel_state + + self.pp_size = parallel_state.get_pipeline_model_parallel_world_size() + self.pp_rank = parallel_state.get_pipeline_model_parallel_rank() + self.is_first_stage = parallel_state.is_pipeline_first_stage() + self.is_last_stage = parallel_state.is_pipeline_last_stage() + except Exception: + # No parallel state initialized (e.g., unit test) + self.pp_size = 1 + self.pp_rank = 0 + self.is_first_stage = True + self.is_last_stage = True + + # Map global layer idx -> local layer idx on this PP stage + self._global_to_local: dict[int, int] = {} + self._local_aux_indices: list[int] = [] + self._local_last_idx: int | None = None + self._compute_local_layer_mapping() + + self._captured: dict[str, Tensor] = {} + self._hooks: list[torch.utils.hooks.RemovableHandle] = [] + + def _compute_local_layer_mapping(self) -> None: + """Map global layer indices to local indices on this PP stage.""" + for local_idx, layer in enumerate(self._decoder.layers): + # Megatron layers have layer_number (1-indexed) + global_idx = int(getattr(layer, "layer_number", local_idx + 1)) - 1 + if global_idx in self.target_layer_ids: + self._global_to_local[global_idx] = local_idx + self._local_aux_indices.append(local_idx) + if global_idx == self.last_layer_idx: + self._local_last_idx = local_idx + + def _make_layer_output_hook(self, key: str): + def hook(_module, _args, output): + # Megatron decoder layer output is typically (hidden_states, ...) + hidden_states = output[0] if isinstance(output, tuple) else output + if hidden_states is None: + return + self._captured[key] = hidden_states.detach().clone() + + return hook + + def _make_embedding_hook(self): + def hook(_module, _args, output): + if isinstance(output, tuple): + output = output[0] + self._captured["embeds"] = output.detach().clone() + + return hook + + def register_hooks(self) -> None: + """Register forward hooks on embedding, target layers, and last layer.""" + self.clear_hooks() + self._captured.clear() + + # Embedding hook (only on first PP stage) + if self.is_first_stage: + embedding = getattr(self.model, "embedding", None) + if embedding is not None: + self._hooks.append(embedding.register_forward_hook(self._make_embedding_hook())) + + # Target layer hooks + for global_idx in self.target_layer_ids: + local_idx = self._global_to_local.get(global_idx) + if local_idx is not None: + layer = self._decoder.layers[local_idx] + self._hooks.append( + layer.register_forward_hook(self._make_layer_output_hook(f"target_layer_{global_idx}")) + ) + + # Last layer hook + if self._local_last_idx is not None: + layer = self._decoder.layers[self._local_last_idx] + self._hooks.append(layer.register_forward_hook(self._make_layer_output_hook("last_layer"))) + + def clear_hooks(self) -> None: + for handle in self._hooks: + handle.remove() + self._hooks.clear() + + @contextmanager + def capture_context(self): + """Context manager that registers hooks on enter and clears on exit.""" + try: + self.register_hooks() + yield self + finally: + self.clear_hooks() + + def _assemble_local_states(self) -> CapturedStates: + """Assemble captured states when PP=1 (all layers on one stage).""" + embeds = self._captured.get("embeds") + + # Concatenate target layer hidden states in target_layer_ids order + hidden_chunks = [] + for global_idx in self.target_layer_ids: + tensor = self._captured.get(f"target_layer_{global_idx}") + if tensor is not None: + hidden_chunks.append(tensor) + + target_hidden = torch.cat(hidden_chunks, dim=-1) if hidden_chunks else None + + last_hidden = self._captured.get("last_layer") + + return CapturedStates( + target_hidden_states=target_hidden, + inputs_embeds=embeds, + target_last_hidden_states=last_hidden, + ) + + @staticmethod + def _send_tensor(tensor: Tensor, dst_rank: int, group) -> None: + """Send a tensor with metadata (shape + dtype).""" + dtype_code = _DTYPE_TO_CODE.get(tensor.dtype) + if dtype_code is None: + raise ValueError(f"Unsupported tensor dtype for send/recv: {tensor.dtype}") + metadata = torch.tensor( + [tensor.shape[0], tensor.shape[1], tensor.shape[2], dtype_code], + dtype=torch.int64, + device=tensor.device, + ) + dist.send(metadata, dst=dst_rank, group=group) + dist.send(tensor.contiguous(), dst=dst_rank, group=group) + + @staticmethod + def _recv_tensor(src_rank: int, group, device: torch.device) -> Tensor: + """Receive a tensor with metadata.""" + metadata = torch.empty(4, dtype=torch.int64, device=device) + dist.recv(metadata, src=src_rank, group=group) + s0, s1, s2, dtype_code = [int(x) for x in metadata.tolist()] + dtype = _CODE_TO_DTYPE.get(dtype_code) + if dtype is None: + raise ValueError(f"Unsupported dtype code in recv: {dtype_code}") + received = torch.empty(s0, s1, s2, dtype=dtype, device=device) + dist.recv(received, src=src_rank, group=group) + return received + + def _gather_distributed(self) -> CapturedStates: + """Gather captured states from all PP stages to the last stage. + + Each target layer's owner rank sends its captured hidden states to the + last PP stage. The embedding is sent from the first stage to the last. + """ + from megatron.core import parallel_state + + pp_group = parallel_state.get_pipeline_model_parallel_group() + last_rank = self.pp_size - 1 + recv_device = torch.device("cuda", torch.cuda.current_device()) + + # If this stage has no captured tensors and is not the last stage, return empty + if not self._captured and not self.is_last_stage: + return CapturedStates() + + gathered_target: dict[int, Tensor] = {} + gathered_last: Tensor | None = None + gathered_embeds: Tensor | None = None + + # Gather target layer hidden states + for global_idx in self.target_layer_ids: + key = f"target_layer_{global_idx}" + tensor = self._captured.get(key) + if tensor is None: + # This stage doesn't own this layer; last stage receives from owner + if self.is_last_stage: + # Determine owner rank (simplified: assume even split) + layers_per_rank = max(1, self.num_layers // self.pp_size) + owner_rank = min(global_idx // layers_per_rank, self.pp_size - 1) + if owner_rank != self.pp_rank: + gathered_target[global_idx] = self._recv_tensor( + src_rank=owner_rank, group=pp_group, device=recv_device + ) + continue + # This stage owns the layer + if self.is_last_stage: + gathered_target[global_idx] = tensor + else: + self._send_tensor(tensor, dst_rank=last_rank, group=pp_group) + + # Gather last layer hidden states + last_tensor = self._captured.get("last_layer") + if last_tensor is not None and not self.is_last_stage: + self._send_tensor(last_tensor, dst_rank=last_rank, group=pp_group) + elif self.is_last_stage and last_tensor is not None: + gathered_last = last_tensor + elif self.is_last_stage: + # Receive from the stage that owns the last layer + layers_per_rank = max(1, self.num_layers // self.pp_size) + owner_rank = min(self.last_layer_idx // layers_per_rank, self.pp_size - 1) + if owner_rank != self.pp_rank: + gathered_last = self._recv_tensor(src_rank=owner_rank, group=pp_group, device=recv_device) + + # Gather embeddings + if self.is_first_stage: + embeds = self._captured.get("embeds") + if embeds is not None: + if self.is_last_stage: + gathered_embeds = embeds + else: + self._send_tensor(embeds, dst_rank=last_rank, group=pp_group) + elif self.is_last_stage: + gathered_embeds = self._recv_tensor(src_rank=0, group=pp_group, device=recv_device) + + if not self.is_last_stage: + return CapturedStates() + + # Concatenate target layers in order + hidden_chunks = [] + for global_idx in self.target_layer_ids: + tensor = gathered_target.get(global_idx) + if tensor is not None: + hidden_chunks.append(tensor) + target_hidden = torch.cat(hidden_chunks, dim=-1) if hidden_chunks else None + + return CapturedStates( + target_hidden_states=target_hidden, + inputs_embeds=gathered_embeds, + target_last_hidden_states=gathered_last, + ) + + def get_captured_states(self) -> CapturedStates: + """Return captured states, gathering across PP stages if needed.""" + if self.pp_size == 1: + return self._assemble_local_states() + return self._gather_distributed() + + +def forward_with_dspark(model, forward_kwargs, batch, target_layer_ids): + from megatron.core import tensor_parallel + from megatron.core.utils import unwrap_model + + capture = HiddenStateCapture(model, target_layer_ids) + with capture.capture_context(): + output_tensor = model(**forward_kwargs) + + captured = capture.get_captured_states() + policy_model = unwrap_model(model) + draft_model = getattr(policy_model, "draft_model", None) + if draft_model is None or captured.target_hidden_states is None: + return output_tensor, None, None + + def batch_first(hidden_states): + if hidden_states is None: + return None + if policy_model.config.sequence_parallel: + hidden_states = tensor_parallel.gather_from_sequence_parallel_region( + hidden_states, tensor_parallel_output_grad=False + ) + return hidden_states.transpose(0, 1).contiguous() + + outputs = draft_model( + input_ids=batch["tokens"], + target_hidden_states=batch_first(captured.target_hidden_states), + loss_mask=batch["full_loss_masks"], + target_last_hidden_states=batch_first(captured.target_last_hidden_states), + ) + return output_tensor, outputs, draft_model.config diff --git a/vime/backends/megatron_utils/dspark/loss.py b/vime/backends/megatron_utils/dspark/loss.py new file mode 100644 index 000000000..b3139c0b3 --- /dev/null +++ b/vime/backends/megatron_utils/dspark/loss.py @@ -0,0 +1,373 @@ +"""DSpark 3-loss computation and DraftLossWrapper. + +Adapted from: + - DeepSpec/deepspec/modeling/dspark/loss.py (3-loss: CE + L1/TV + Confidence) + - NeMo RL/nemo_rl/algorithms/loss/wrapper.py (DraftLossWrapper pattern) + +The 3 losses are: + L_ce: cross-entropy of draft logits vs target token ids (weight 0.1) + L_tv: L1 distance between draft probs and target probs (weight 0.9) + L_conf: BCE of confidence pred vs empirical accept rate (weight 1.0) + +All losses use position decay: weight *= exp(-pos / gamma), gamma=4.0. + +Loss denominators are all-reduced across the data-parallel group to normalize +correctly. The final backward loss is scaled by world_size to counteract +DDP's gradient averaging. +""" + +import logging + +import torch +import torch.distributed as dist +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint + +from .common import DSparkConfig, DSparkForwardOutput + +logger = logging.getLogger(__name__) + +# Chunk size for vocab-dimension operations. Each chunk creates a temporary +# tensor of shape (batch, seq, block_size, chunk_size). With chunk_size=16384 +# and typical batch*seq*block_size=~3500, each chunk is ~235 MB (float32). +_L1_VOCAB_CHUNK_SIZE = 16384 + + +def _all_reduce_loss_denominators( + loss_terms: dict[str, torch.Tensor], + *, + world_size: int, +) -> dict[str, torch.Tensor]: + """All-reduce loss denominators across DP group for normalization.""" + denominators = {} + for key in ("ce_loss_den", "l1_loss_den", "confidence_loss_den"): + tensor = loss_terms[key].detach().clone() + if world_size > 1: + try: + from megatron.core import mpu + + dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + if dp_group is not None: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=dp_group) + else: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + except Exception: + # Fallback: global all_reduce + if dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + denominators[key] = tensor + return denominators + + +def _build_loss_weight_mask( + *, + eval_mask: torch.Tensor, + block_size: int, + device: torch.device, + loss_decay_gamma: float | None, +) -> torch.Tensor: + """Build per-position loss weight with exponential decay.""" + loss_weight_mask = eval_mask.to(torch.float32) + if loss_decay_gamma is not None and loss_decay_gamma > 0: + positions = torch.arange(block_size, device=device).view(1, 1, -1) + decay_weights = torch.exp(-positions.float() / float(loss_decay_gamma)) + loss_weight_mask = loss_weight_mask * decay_weights + return loss_weight_mask + + +def _compute_accept_rate_3d( + *, + outputs: DSparkForwardOutput, + aligned_target_logits: torch.Tensor | None, +) -> torch.Tensor | None: + """Compute per-position acceptance rate: 1 - 0.5 * L1(draft_probs, target_probs). + + Computed without gradients (only used as a detached target for the + confidence head). Uses chunked logsumexp to avoid materializing the full + vocab-dimension softmax or difference tensors. + """ + if aligned_target_logits is None: + return None + with torch.no_grad(): + draft_logits = outputs.draft_logits.float() + target_logits = aligned_target_logits.float() + vocab_size = draft_logits.shape[-1] + log_Z_draft = torch.logsumexp(draft_logits, dim=-1, keepdim=True) + log_Z_target = torch.logsumexp(target_logits, dim=-1, keepdim=True) + l1_shape = draft_logits.shape[:-1] + l1_dist = torch.zeros(l1_shape, device=draft_logits.device, dtype=torch.float32) + for start in range(0, vocab_size, _L1_VOCAB_CHUNK_SIZE): + end = min(start + _L1_VOCAB_CHUNK_SIZE, vocab_size) + pa = torch.exp(draft_logits[..., start:end] - log_Z_draft) + pb = torch.exp(target_logits[..., start:end] - log_Z_target) + l1_dist += (pa - pb).abs().sum(dim=-1) + accept_rate_3d = (1.0 - 0.5 * l1_dist).clamp_(0.0, 1.0) + return accept_rate_3d + + +def _compute_local_l1_term( + *, + outputs: DSparkForwardOutput, + aligned_target_logits: torch.Tensor | None, + loss_weight_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute L1/TV loss numerator and denominator. + + Uses logsumexp + chunked exp + gradient checkpointing to avoid OOM. + The full softmax tensors are never materialized; each vocab chunk is + recomputed during backward via torch.utils.checkpoint. + """ + zero = outputs.draft_logits.new_zeros((), dtype=torch.float32) + if aligned_target_logits is None: + return zero, zero + draft_logits = outputs.draft_logits.float() + target_logits = aligned_target_logits.float() + vocab_size = draft_logits.shape[-1] + + # Log partition functions (small, no OOM risk) + log_Z_draft = torch.logsumexp(draft_logits, dim=-1, keepdim=True) + log_Z_target = torch.logsumexp(target_logits, dim=-1, keepdim=True) + + # Chunked L1: sum_i |exp(a_i - log_Z_a) - exp(b_i - log_Z_b)| + # Each chunk creates a temporary of shape (..., chunk_size), well under + # the torch_memory_saver margin. Gradient checkpointing ensures chunk + # intermediates are NOT saved for backward (recomputed instead). + l1_shape = draft_logits.shape[:-1] + l1_dist = torch.zeros(l1_shape, device=draft_logits.device, dtype=torch.float32) + + def _chunk_l1(a_chunk, b_chunk, log_za, log_zb): + pa = torch.exp(a_chunk - log_za) + pb = torch.exp(b_chunk - log_zb) + return (pa - pb).abs().sum(dim=-1) + + for start in range(0, vocab_size, _L1_VOCAB_CHUNK_SIZE): + end = min(start + _L1_VOCAB_CHUNK_SIZE, vocab_size) + a_slice = draft_logits[..., start:end] + b_slice = target_logits[..., start:end] + chunk_l1 = checkpoint.checkpoint( + _chunk_l1, + a_slice, + b_slice, + log_Z_draft, + log_Z_target, + use_reentrant=False, + ) + l1_dist = l1_dist + chunk_l1 + + l1_loss_num = (l1_dist * loss_weight_mask).sum() + l1_loss_den = loss_weight_mask.sum() + return l1_loss_num, l1_loss_den + + +def _collect_local_terms( + *, + outputs: DSparkForwardOutput, + loss_decay_gamma: float | None, + l1_loss_alpha: float, +) -> tuple[dict[str, torch.Tensor], bool]: + """Collect local loss terms (numerators and denominators).""" + draft_logits = outputs.draft_logits + target_ids = outputs.target_ids + eval_mask = outputs.eval_mask + _, _, block_size, vocab_size = draft_logits.shape + device = draft_logits.device + + loss_weight_mask = _build_loss_weight_mask( + eval_mask=eval_mask, + block_size=block_size, + device=device, + loss_decay_gamma=loss_decay_gamma, + ) + flat_logits = draft_logits.reshape(-1, vocab_size) + flat_targets = target_ids.reshape(-1) + flat_weights = loss_weight_mask.reshape(-1) + loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") + ce_loss_num = (loss_per_token * flat_weights).sum() + ce_loss_den = flat_weights.sum() + + aligned_target_logits = outputs.aligned_target_logits + accept_rate_3d = _compute_accept_rate_3d( + outputs=outputs, + aligned_target_logits=aligned_target_logits, + ) + zero = ce_loss_num.new_zeros(()) + + assert ( + l1_loss_alpha <= 0 or aligned_target_logits is not None + ), "aligned_target_logits is required when l1_loss_alpha > 0." + if l1_loss_alpha > 0: + l1_loss_num, l1_loss_den = _compute_local_l1_term( + outputs=outputs, + aligned_target_logits=aligned_target_logits, + loss_weight_mask=loss_weight_mask, + ) + else: + l1_loss_num = zero + l1_loss_den = zero + + has_confidence = outputs.confidence_pred is not None + confidence_loss_num = zero + confidence_loss_den = zero + if has_confidence: + assert accept_rate_3d is not None, "aligned_target_logits is required when confidence head is enabled." + confidence_targets = accept_rate_3d.detach() + confidence_errors = ( + F.binary_cross_entropy_with_logits( + outputs.confidence_pred.float(), + confidence_targets, + reduction="none", + ) + * loss_weight_mask + ) + confidence_loss_num = confidence_errors.sum() + confidence_loss_den = loss_weight_mask.sum() + + loss_terms = { + "ce_loss_num": ce_loss_num, + "ce_loss_den": ce_loss_den, + "l1_loss_num": l1_loss_num, + "l1_loss_den": l1_loss_den, + "confidence_loss_num": confidence_loss_num, + "confidence_loss_den": confidence_loss_den, + } + return loss_terms, has_confidence + + +def _build_loss( + *, + loss_terms: dict[str, torch.Tensor], + global_denominators: dict[str, torch.Tensor], + ce_loss_alpha: float, + l1_loss_alpha: float, + confidence_head_alpha: float, + has_confidence: bool, + world_size: int, +) -> torch.Tensor: + """Build the final backward loss from local terms and global denominators.""" + ce_loss = loss_terms["ce_loss_num"] / (global_denominators["ce_loss_den"] + 1e-6) + l1_loss = ce_loss.new_zeros(()) + if global_denominators["l1_loss_den"].item() > 0: + l1_loss = loss_terms["l1_loss_num"] / (global_denominators["l1_loss_den"] + 1e-6) + confidence_loss = ce_loss.new_zeros(()) + if has_confidence: + confidence_loss = loss_terms["confidence_loss_num"] / (global_denominators["confidence_loss_den"] + 1e-6) + return (ce_loss_alpha * ce_loss + l1_loss_alpha * l1_loss + confidence_head_alpha * confidence_loss) * world_size + + +def compute_dspark_loss( + *, + outputs: DSparkForwardOutput, + config: DSparkConfig, +) -> tuple[torch.Tensor, dict[str, float]]: + """Compute DSpark 3-loss. + + Args: + outputs: DSparkForwardOutput from DSparkModel.forward + config: DSparkConfig with loss weights and decay gamma + Returns: + (backward_loss, metrics_dict) + - backward_loss: scalar tensor for backward() + - metrics_dict: dict of float values for logging + """ + loss_terms, has_confidence = _collect_local_terms( + outputs=outputs, + loss_decay_gamma=config.loss_decay_gamma, + l1_loss_alpha=float(config.l1_loss_alpha), + ) + + world_size = dist.get_world_size() if dist.is_initialized() else 1 + global_denominators = _all_reduce_loss_denominators(loss_terms, world_size=world_size) + + # Local loss for logging + local_ce_loss = loss_terms["ce_loss_num"] / (loss_terms["ce_loss_den"] + 1e-6) + local_l1_loss = local_ce_loss.new_zeros(()) + if loss_terms["l1_loss_den"].item() > 0: + local_l1_loss = loss_terms["l1_loss_num"] / (loss_terms["l1_loss_den"] + 1e-6) + local_confidence_loss = local_ce_loss.new_zeros(()) + if has_confidence: + local_confidence_loss = loss_terms["confidence_loss_num"] / (loss_terms["confidence_loss_den"] + 1e-6) + + backward_loss = _build_loss( + loss_terms=loss_terms, + global_denominators=global_denominators, + ce_loss_alpha=float(config.ce_loss_alpha), + l1_loss_alpha=float(config.l1_loss_alpha), + confidence_head_alpha=float(config.confidence_head_alpha), + has_confidence=has_confidence, + world_size=world_size, + ) + + metrics = { + "dspark/ce_loss": float(local_ce_loss.detach().item()), + "dspark/l1_loss": float(local_l1_loss.detach().item()), + "dspark/confidence_loss": float(local_confidence_loss.detach().item()), + "dspark/total_loss": float( + ( + config.ce_loss_alpha * local_ce_loss + + config.l1_loss_alpha * local_l1_loss + + config.confidence_head_alpha * local_confidence_loss + ) + .detach() + .item() + ), + } + return backward_loss, metrics + + +class DraftLossWrapper: + """Combine policy RL loss with DSpark draft loss. + + Following NeMo RL's DraftLossWrapper pattern: + combined_loss = policy_loss + draft_loss_weight * draft_loss + + The draft loss is computed from DSparkForwardOutput and added to the + policy loss. The policy loss function is called first, then the draft + loss is computed and combined. + """ + + def __init__(self, config: DSparkConfig): + self.config = config + self.draft_loss_weight = config.draft_loss_weight + + def __call__( + self, + policy_loss: torch.Tensor, + dspark_outputs: DSparkForwardOutput, + ) -> tuple[torch.Tensor, dict[str, float]]: + """Compute combined loss. + + Args: + policy_loss: scalar tensor, the RL policy loss (already computed) + dspark_outputs: DSparkForwardOutput from DSparkModel.forward + Returns: + (combined_loss, dspark_metrics) + """ + draft_loss, dspark_metrics = compute_dspark_loss( + outputs=dspark_outputs, + config=self.config, + ) + combined_loss = policy_loss + self.draft_loss_weight * draft_loss + return combined_loss, dspark_metrics + + +def build_combined_loss_fn(policy_loss_fn, args, batch, num_microbatches, global_batch_size, outputs, config): + wrapper = DraftLossWrapper(config) + + def combined_loss_fn(logits): + if args.dspark_freeze_policy: + logits = logits.detach() + policy_loss, num_elems, metrics = policy_loss_fn( + args, + batch, + num_microbatches, + global_batch_size, + logits, + ) + combined_loss, dspark_metrics = wrapper(policy_loss, outputs) + keys = list(dspark_metrics) + values = metrics["values"].new_tensor([dspark_metrics[key] for key in keys]) + metrics["keys"] += keys + metrics["values"] = torch.cat((metrics["values"], values)) + return combined_loss, num_elems, metrics + + return combined_loss_fn diff --git a/vime/backends/megatron_utils/dspark/markov_head.py b/vime/backends/megatron_utils/dspark/markov_head.py new file mode 100644 index 000000000..edf33ad45 --- /dev/null +++ b/vime/backends/megatron_utils/dspark/markov_head.py @@ -0,0 +1,219 @@ +"""DSpark Markov head: sequential bias applied to parallel backbone logits. + +Adapted from DeepSpec/deepspec/modeling/dspark/markov_head.py. + +The Markov head provides a per-position bias to the draft logits, enabling +semi-autoregressive generation: the backbone produces base logits for all +positions in parallel, then the Markov head adds a bias that depends on the +previous token (teacher-forced during training, autoregressive at inference). + +Three variants are supported: + - vanilla: W2(W1[x_{k-1}]) + - gated: W2(gate * W1[x_{k-1}]), gate = sigmoid(Linear([h, W1[x]])) + - rnn: GRU-like recurrent state carrying prefix history + +For vime online training, we only need the training forward (``apply_block_logits``); +sampling methods (``sample_block_tokens``) are used at inference by vLLM, not here. +""" + +import torch +from torch import nn + + +class VanillaMarkov(nn.Module): + """Vanilla Markov head: bias = W2(W1[prev_token]).""" + + def __init__(self, *, vocab_size: int, markov_rank: int): + super().__init__() + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + self.markov_head_type = "vanilla" + assert self.markov_rank > 0, f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." + self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) + self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) + + def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w1(token_ids.long()) + + def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor: + return self.markov_w2(latent_states) + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + del hidden_states + return self.project_bias(self.get_prev_embeddings(token_ids)) + + def apply_step_logits( + self, + logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + return logits + self.compute_step_bias(token_ids, hidden_states) + + def apply_block_logits( + self, + base_logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + """Apply Markov bias to all positions in a block (teacher-forced). + + Args: + base_logits: [B, num_blocks, block_size, V] + token_ids: [B, num_blocks, block_size] (prev token per position) + hidden_states: unused for vanilla + Returns: + [B, num_blocks, block_size, V] + """ + if base_logits.size(2) == 0: + return base_logits + markov_bias = self.compute_step_bias(token_ids, hidden_states) + return base_logits + markov_bias + + +class GatedMarkovHead(VanillaMarkov): + """Gated Markov head: bias = W2(gate * W1[prev_token]), gate from [h, W1[x]].""" + + def __init__( + self, + *, + vocab_size: int, + markov_rank: int, + hidden_size: int, + ): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.markov_head_type = "gated" + self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank) + + def compute_gate( + self, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + assert hidden_states is not None + prev_embeddings = self.get_prev_embeddings(token_ids) + gate_inputs = torch.cat([hidden_states, prev_embeddings], dim=-1) + return torch.sigmoid(self.gate_proj(gate_inputs)) + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + prev_embeddings = self.get_prev_embeddings(token_ids) + gate = self.compute_gate(token_ids, hidden_states).to(dtype=prev_embeddings.dtype) + return self.project_bias(gate * prev_embeddings) + + +class RNNHead(VanillaMarkov): + """RNN-based head with GRU-like recurrent state across positions in a block.""" + + def __init__( + self, + *, + vocab_size: int, + markov_rank: int, + hidden_size: int, + ): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.markov_head_type = "rnn" + self.hidden_size = hidden_size + self.state_size = markov_rank + self.joint_proj = nn.Linear(2 * markov_rank + hidden_size, 3 * markov_rank) + + def _rnn_step( + self, + state: torch.Tensor, + prev_embeddings: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) + proj = self.joint_proj(z) + gate_raw, candidate_raw, output_raw = proj.chunk(3, dim=-1) + gate = torch.sigmoid(gate_raw) + candidate = torch.tanh(candidate_raw) + new_state = gate * state + (1.0 - gate) * candidate + bias = self.project_bias(torch.tanh(output_raw)) + return new_state, bias + + def compute_step_bias( + self, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + """Stateless single-step bias (state initialized to zero).""" + assert hidden_states is not None + prev_embeddings = self.get_prev_embeddings(token_ids) + state = torch.zeros_like(prev_embeddings) + _, bias = self._rnn_step(state, prev_embeddings, hidden_states) + return bias + + def apply_block_logits( + self, + base_logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + """Apply RNN bias during training (teacher-forced, unrolled over block_size). + + Args: + base_logits: [B, num_blocks, block_size, V] + token_ids: [B, num_blocks, block_size] + hidden_states: [B, num_blocks, block_size, d] + """ + assert hidden_states is not None + block_size = base_logits.size(-2) + if block_size == 0: + return base_logits + + leading_shape = base_logits.shape[:-2] + state = torch.zeros( + *leading_shape, + self.markov_rank, + device=base_logits.device, + dtype=hidden_states.dtype, + ) + + output_logits = [] + for k in range(block_size): + prev_emb = self.get_prev_embeddings(token_ids[..., k]) + h_k = hidden_states[..., k, :] + state, bias = self._rnn_step(state, prev_emb, h_k) + output_logits.append(base_logits[..., k, :] + bias) + + return torch.stack(output_logits, dim=-2) + + +def build_markov_head(config) -> nn.Module | None: + """Build a Markov head from a DSparkConfig (or compatible object).""" + markov_rank = int(config.markov_rank) + assert markov_rank >= 0, f"markov_rank must be >= 0, got {markov_rank}" + if markov_rank == 0: + return None + + markov_head_type = str(config.markov_head_type).lower() + if markov_head_type == "vanilla": + return VanillaMarkov( + vocab_size=config.vocab_size, + markov_rank=markov_rank, + ) + if markov_head_type == "gated": + return GatedMarkovHead( + vocab_size=config.vocab_size, + markov_rank=markov_rank, + hidden_size=config.hidden_size, + ) + if markov_head_type == "rnn": + return RNNHead( + vocab_size=config.vocab_size, + markov_rank=markov_rank, + hidden_size=config.hidden_size, + ) + raise ValueError(f"Unsupported markov_head_type: {markov_head_type!r}") diff --git a/vime/backends/megatron_utils/dspark/modeling.py b/vime/backends/megatron_utils/dspark/modeling.py new file mode 100644 index 000000000..1f197e768 --- /dev/null +++ b/vime/backends/megatron_utils/dspark/modeling.py @@ -0,0 +1,602 @@ +"""DSpark draft model for vime Megatron backend. + +Adapted from DeepSpec/deepspec/modeling/dspark/qwen3/modeling.py:Qwen3DSparkModel. + +The DSpark draft model is a semi-autoregressive speculative decoding drafter: + - Parallel backbone (5 decoder layers) produces draft hidden states + - Markov head adds sequential bias to draft logits + - Confidence head predicts per-position acceptance probability + +Unlike Eagle3 (which uses ModelOpt's EagleModule with Megatron TP support), +DSpark's dual-input attention has no existing Megatron implementation, so +this module is replicated on every TP rank and uses plain ``nn.Linear`` + SDPA. + +The model is attached as ``policy_chunk.draft_model`` before DDP wrapping, +following the NeMo RL Eagle3 pattern. +""" + +import logging + +import torch +import torch.nn.functional as F +from torch import nn + +from .attention import DSparkParallelAttention +from .common import ( + AcceptRatePredictor, + DSparkConfig, + DSparkForwardOutput, + build_eval_mask, + create_dspark_attention_mask, + create_noise_embed, + create_position_ids, + sample_anchor_positions, +) +from .markov_head import build_markov_head + +logger = logging.getLogger(__name__) + + +def _all_gather_vocab_weight(weight: torch.Tensor) -> torch.Tensor: + """All-gather a TP-sharded vocab weight along dim=0. + + Megatron's VocabParallelEmbedding shards the vocab dimension across TP + ranks. Each rank holds [vocab_size // TP, hidden_size]. This function + collects the full [padded_vocab_size, hidden_size] tensor on every rank. + + Returns the original weight if all-gather fails (e.g., TP group not + accessible); callers should handle shape mismatch as a fallback. + """ + import torch.distributed as dist + + if not dist.is_initialized(): + return weight + + tp_group = None + for module_path, func_name in [ + ("megatron.core.tensor_parallel", "get_tensor_model_parallel_group"), + ("megatron.core.parallel_state", "get_tensor_model_parallel_group"), + ]: + try: + module = __import__(module_path, fromlist=[func_name]) + tp_group = getattr(module, func_name)() + break + except (ImportError, AttributeError, RuntimeError): + continue + + if tp_group is None: + return weight + + try: + world_size = dist.get_world_size(group=tp_group) + except (RuntimeError, ValueError): + return weight + + if world_size <= 1: + return weight + + tensor_list = [torch.empty_like(weight) for _ in range(world_size)] + dist.all_gather(tensor_list, weight.contiguous().detach(), group=tp_group) + return torch.cat(tensor_list, dim=0).detach() + + +class DSparkRMSNorm(nn.Module): + """RMSNorm matching Qwen3's normalization.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states.to(input_dtype) + + +class DSparkMLP(nn.Module): + """SwiGLU MLP matching Qwen3.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class DSparkDecoderLayer(nn.Module): + """One DSpark decoder layer with dual-input attention + SwiGLU MLP.""" + + def __init__(self, config: DSparkConfig): + super().__init__() + self.self_attn = DSparkParallelAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + attention_bias=False, + rms_norm_eps=config.rms_norm_eps, + rotary_base=config.rotary_base, + ) + self.mlp = DSparkMLP( + hidden_size=config.hidden_size, + intermediate_size=_compute_intermediate_size(config), + ) + self.input_layernorm = DSparkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DSparkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + target_hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + target_hidden_states=target_hidden_states, + position_ids=position_ids, + attention_mask=attention_mask, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +def _compute_intermediate_size(config: DSparkConfig) -> int: + """Compute MLP intermediate size from config. + + For Qwen3 models, intermediate_size is typically ~2.75x hidden_size. + This can be overridden by setting ``intermediate_size`` on the config. + """ + if hasattr(config, "intermediate_size") and config.intermediate_size > 0: + return config.intermediate_size + # Qwen3-4B/8B: hidden=2560/4096, intermediate=6912/12288 + # Ratio ~2.7. Use a multiple of 256 for efficiency. + raw = int(config.hidden_size * 2.75) + return ((raw + 255) // 256) * 256 + + +class DSparkModel(nn.Module): + """DSpark draft model: parallel backbone + Markov head + confidence head. + + This is a plain ``nn.Module`` attached as ``policy_chunk.draft_model`` and + replicated on every TP rank. DDP wrapping on the parent policy chunk still + covers its parameters. + + Structure: + - embed_tokens: shared from policy (frozen by default) + - layers: ``num_draft_layers`` DSparkDecoderLayer + - norm: final RMSNorm + - fc: projection from [num_target_layers * hidden] -> hidden + - hidden_norm: RMSNorm on target hidden states before fc + - lm_head: shared from policy (frozen by default) + - markov_head: vanilla/gated/rnn (default vanilla, rank=256) + - confidence_head: AcceptRatePredictor + """ + + def __init__(self, config: DSparkConfig): + super().__init__() + self.config = config + self.target_layer_ids = list(config.target_layer_ids) + self.block_size = config.block_size + self.mask_token_id = config.mask_token_id + self.num_anchors = config.num_anchors + + # Embedding and LM head — will be shared from policy via initialize_embeddings_and_head + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Backbone + self.layers = nn.ModuleList([DSparkDecoderLayer(config) for _ in range(config.num_draft_layers)]) + self.norm = DSparkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # FC projection: [num_target_layers * hidden] -> hidden + self.fc = nn.Linear( + len(self.target_layer_ids) * config.hidden_size, + config.hidden_size, + bias=False, + ) + self.hidden_norm = DSparkRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Markov head + self.markov_head = build_markov_head(config) + + # Confidence head + self.enable_confidence_head = config.enable_confidence_head + self.confidence_head_with_markov = config.confidence_head_with_markov + self.confidence_head = None + if self.enable_confidence_head: + input_dim = config.hidden_size + if self.confidence_head_with_markov: + input_dim += config.markov_rank + self.confidence_head = AcceptRatePredictor(input_dim=input_dim) + + def initialize_embeddings_and_head( + self, + *, + embed_tokens: nn.Module, + lm_head: nn.Module, + freeze: bool = True, + ): + """Copy policy's embedding and lm_head weights to DSpark (shared). + + When policy uses TP>1, Megatron's VocabParallelEmbedding shards + the vocab dimension across TP ranks. We all-gather to reconstruct + the full vocab embedding for the replicated DSpark model. + """ + embed_weight = embed_tokens.weight.detach() + lm_head_weight = lm_head.weight.detach() + + # All-gather if policy embedding is TP-sharded + if embed_weight.shape != self.embed_tokens.weight.shape: + embed_weight = _all_gather_vocab_weight(embed_weight) + if lm_head_weight.shape != self.lm_head.weight.shape: + lm_head_weight = _all_gather_vocab_weight(lm_head_weight) + + with torch.no_grad(): + # If all-gather succeeded, shapes match and we copy directly. + # If all-gather failed (TP group not accessible), copy only the + # matching portion; the rest stays zero-initialized and will be + # overwritten by load_pretrained_weights() if a pretrained + # checkpoint is provided. + if embed_weight.shape == self.embed_tokens.weight.shape: + self.embed_tokens.weight.copy_(embed_weight) + else: + min_rows = min(embed_weight.shape[0], self.embed_tokens.weight.shape[0]) + logger.warning( + "[DSpark] embed_tokens shape mismatch after all-gather: " + "dspark %s vs policy %s, copying first %d rows", + self.embed_tokens.weight.shape, + embed_weight.shape, + min_rows, + ) + self.embed_tokens.weight.zero_() + self.embed_tokens.weight[:min_rows].copy_(embed_weight[:min_rows]) + + if lm_head_weight.shape == self.lm_head.weight.shape: + self.lm_head.weight.copy_(lm_head_weight) + else: + min_rows = min(lm_head_weight.shape[0], self.lm_head.weight.shape[0]) + logger.warning( + "[DSpark] lm_head shape mismatch after all-gather: " + "dspark %s vs policy %s, copying first %d rows", + self.lm_head.weight.shape, + lm_head_weight.shape, + min_rows, + ) + self.lm_head.weight.zero_() + self.lm_head.weight[:min_rows].copy_(lm_head_weight[:min_rows]) + + if freeze: + self.set_embedding_head_trainable(False) + + def load_pretrained_weights(self, path: str): + """Load pre-trained DSpark weights from safetensors file. + + Loads ALL parameters including embed_tokens and lm_head. + The pre-trained DSpark checkpoint has untied embeddings + (tie_word_embeddings=false), so its lm_head is very different + from the policy's tied lm_head. Using the policy's lm_head + would produce random predictions (CE loss ~11.0). + """ + from safetensors.torch import load_file + + state_dict = load_file(path) + loaded = 0 + missing = 0 + mismatched = 0 + for name, param in self.named_parameters(): + if name in state_dict: + tensor = state_dict[name] + if tensor.shape == param.data.shape: + param.data.copy_(tensor.to(param.dtype)) + loaded += 1 + elif ( + tensor.dim() == param.data.dim() + and tensor.shape[1:] == param.data.shape[1:] + and tensor.shape[0] <= param.data.shape[0] + ): + # Vocab padding: checkpoint has fewer rows (original vocab), + # model has padded vocab for TP. Zero-pad and copy. + with torch.no_grad(): + param.data.zero_() + param.data[: tensor.shape[0]].copy_(tensor.to(param.dtype)) + loaded += 1 + logger.info( + "[DSpark] Padded %s: checkpoint %s -> model %s (vocab padding)", + name, + tuple(tensor.shape), + tuple(param.data.shape), + ) + else: + logger.warning( + "[DSpark] Shape mismatch for %s: checkpoint %s vs model %s", + name, + tuple(tensor.shape), + tuple(param.data.shape), + ) + mismatched += 1 + else: + logger.warning("[DSpark] Missing key in pretrained: %s", name) + missing += 1 + total = sum(1 for _, _ in self.named_parameters()) + logger.info( + "[DSpark] Loaded %d/%d params from pretrained checkpoint " "(missing=%d, mismatched=%d)", + loaded, + total, + missing, + mismatched, + ) + + def set_embedding_head_trainable(self, trainable: bool): + self.embed_tokens.requires_grad_(trainable) + self.lm_head.requires_grad_(trainable) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) + + def predict_confidence_step( + self, + hidden_states: torch.Tensor, + prev_token_ids: torch.Tensor | None = None, + ) -> torch.Tensor | None: + """Compute confidence predictions for each position. + + Args: + hidden_states: [bsz, num_blocks, block_size, hidden] + prev_token_ids: [bsz, num_blocks, block_size] (needed if confidence_head_with_markov) + Returns: + [bsz, num_blocks, block_size] or None + """ + if self.confidence_head is None: + return None + if self.confidence_head_with_markov: + assert self.markov_head is not None + assert prev_token_ids is not None + prev_embeddings = self.markov_head.get_prev_embeddings(prev_token_ids).to(dtype=hidden_states.dtype) + features = torch.cat([hidden_states, prev_embeddings], dim=-1) + return self.confidence_head(features).float() + return self.confidence_head(hidden_states).float() + + def forward( + self, + input_ids: torch.Tensor, + target_hidden_states: torch.Tensor, + loss_mask: torch.Tensor, + target_last_hidden_states: torch.Tensor | None = None, + ) -> DSparkForwardOutput: + """DSpark training forward. + + Args: + input_ids: [bsz, seq_len] token ids from policy + target_hidden_states: [bsz, seq_len, num_target_layers * hidden] + Concatenated hidden states from policy's target_layer_ids. + loss_mask: [bsz, seq_len] loss mask (1 for supervised tokens) + target_last_hidden_states: [bsz, seq_len, hidden] (optional) + Policy's last layer hidden states, for L_tv/L_conf computation. + Returns: + DSparkForwardOutput + """ + bsz, seq_len = input_ids.shape + device = input_ids.device + + # 1. Sample anchor positions + anchor_positions, block_keep_mask = sample_anchor_positions( + seq_len=seq_len, + loss_mask=loss_mask, + num_anchors=self.num_anchors, + device=device, + ) + + # 2. Create noise embedding (draft input) + noise_embedding = create_noise_embed( + self.embed_tokens, + input_ids, + anchor_positions, + block_keep_mask, + mask_token_id=self.mask_token_id, + block_size=self.block_size, + ) + + # 3. Position ids: context + draft + context_position_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + draft_position_ids = create_position_ids(anchor_positions, self.block_size) + full_position_ids = torch.cat([context_position_ids, draft_position_ids], dim=1) + + # 4. Project target hidden states + # Detach: prevent DSpark loss gradient from flowing back through the + # policy model. The draft model learns to predict the target, not + # the other way around. + target_hidden_projected = self.hidden_norm(self.fc(target_hidden_states.detach())) + + # 5. Build attention mask + attn_mask = create_dspark_attention_mask( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + seq_len=seq_len, + block_size=self.block_size, + device=device, + dtype=noise_embedding.dtype, + ) # [bsz, 1, q_len, ctx_len + q_len] + + # 6. Backbone forward + hidden_states = noise_embedding # [bsz, q_len, hidden] + # Position ids for rotary: need [bsz, ctx_len + q_len] + # The rotary is applied to K which spans ctx_len + q_len + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + target_hidden_states=target_hidden_projected, + position_ids=full_position_ids, + attention_mask=attn_mask, + ) + output_hidden = self.norm(hidden_states) # [bsz, q_len, hidden] + + # 7. Reshape to [bsz, num_blocks, block_size, hidden] + num_blocks = anchor_positions.size(1) + output_hidden_4d = output_hidden.reshape(bsz, num_blocks, self.block_size, -1) + + # 8. Compute target ids (labels for CE loss) + label_offsets = torch.arange(1, self.block_size + 1, device=device).view(1, 1, -1) # [1, 1, block_size] + label_indices = anchor_positions.unsqueeze(-1) + label_offsets # [bsz, num_blocks, block_size] + safe_label_indices = label_indices.clamp(max=seq_len - 1) + safe_label_indices = torch.where( + block_keep_mask.unsqueeze(-1), + safe_label_indices, + torch.zeros_like(safe_label_indices), + ) + target_ids = torch.gather( + input_ids.unsqueeze(1).expand(-1, anchor_positions.size(1), -1), + 2, + safe_label_indices, + ) # [bsz, num_blocks, block_size] + + # 9. Compute aligned target logits (for L_tv / L_conf) + aligned_target_logits = None + if target_last_hidden_states is not None: + target_pred_indices = (safe_label_indices - 1).clamp(min=0) + aligned_target_hidden = torch.gather( + target_last_hidden_states.unsqueeze(1).expand(-1, anchor_positions.size(1), -1, -1), + 2, + target_pred_indices.unsqueeze(-1).expand(-1, -1, -1, target_last_hidden_states.size(-1)), + ) # [bsz, num_blocks, block_size, hidden] + # Detach target logits: the l1_loss gradient must NOT flow + # back through the policy model. The draft model is trained to + # predict the target, not the other way around. + aligned_target_logits = self.compute_logits(aligned_target_hidden).detach() + + # 10. Build eval mask + eval_mask = build_eval_mask( + seq_len=seq_len, + loss_mask=loss_mask, + label_indices=label_indices, + safe_label_indices=safe_label_indices, + block_keep_mask=block_keep_mask, + ) # [bsz, num_blocks, block_size] bool + + # 11. Compute draft logits + anchor_token_ids = torch.gather(input_ids, 1, anchor_positions) # [bsz, num_blocks] + prev_token_ids = torch.cat( + [anchor_token_ids.unsqueeze(-1), target_ids[:, :, :-1]], + dim=-1, + ) # [bsz, num_blocks, block_size] + + draft_logits = self.compute_logits(output_hidden).reshape( + bsz, num_blocks, self.block_size, -1 + ) # [bsz, num_blocks, block_size, vocab] + + # 12. Apply Markov head bias + if self.markov_head is not None: + draft_logits = self.markov_head.apply_block_logits( + draft_logits, + token_ids=prev_token_ids, + hidden_states=output_hidden_4d, + ) + + # 13. Confidence prediction + confidence_pred = None + if self.confidence_head is not None: + confidence_pred = self.predict_confidence_step( + output_hidden_4d, prev_token_ids + ) # [bsz, num_blocks, block_size] + + return DSparkForwardOutput( + draft_logits=draft_logits, + target_ids=target_ids, + eval_mask=eval_mask, + block_keep_mask=block_keep_mask, + confidence_pred=confidence_pred, + aligned_target_logits=aligned_target_logits, + ) + + +def build_dspark_model( + dspark_config: DSparkConfig, + policy_embed_tokens: nn.Module, + policy_lm_head: nn.Module, + pretrained_model_path: str | None = None, +) -> DSparkModel: + """Build a DSpark draft model and initialize shared weights from policy. + + Args: + dspark_config: DSparkConfig with model dims populated + policy_embed_tokens: policy's embedding layer (for weight sharing) + policy_lm_head: policy's lm_head layer (for weight sharing) + pretrained_model_path: Optional path to pre-trained DSpark safetensors + file. If provided, backbone weights are loaded from this file + instead of random init. embed_tokens/lm_head are always copied + from the policy model. + Returns: + DSparkModel (not yet attached to policy chunk) + """ + model = DSparkModel(dspark_config) + model.initialize_embeddings_and_head( + embed_tokens=policy_embed_tokens, + lm_head=policy_lm_head, + freeze=True, + ) + if pretrained_model_path is not None: + model.load_pretrained_weights(pretrained_model_path) + logger.info( + "[DSpark] Built DSpark draft model: %d layers, block_size=%d, " + "target_layer_ids=%s, markov_rank=%d, confidence=%s, " + "intermediate_size=%d, pretrained=%s", + dspark_config.num_draft_layers, + dspark_config.block_size, + dspark_config.target_layer_ids, + dspark_config.markov_rank, + dspark_config.enable_confidence_head, + _compute_intermediate_size(dspark_config), + pretrained_model_path is not None, + ) + return model + + +def attach_dspark_model(model, args, config) -> None: + dspark_config = DSparkConfig( + block_size=args.dspark_block_size, + num_draft_layers=args.dspark_num_draft_layers, + target_layer_ids=args.dspark_target_layer_ids, + mask_token_id=args.dspark_mask_token_id, + num_anchors=args.dspark_num_anchors, + markov_rank=args.dspark_markov_rank, + markov_head_type=args.dspark_markov_head_type, + enable_confidence_head=not args.dspark_disable_confidence_head, + ce_loss_alpha=args.dspark_ce_loss_alpha, + l1_loss_alpha=args.dspark_l1_loss_alpha, + confidence_head_alpha=args.dspark_confidence_head_alpha, + loss_decay_gamma=args.dspark_loss_decay_gamma, + draft_loss_weight=args.dspark_draft_loss_weight, + hidden_size=config.hidden_size, + vocab_size=args.padded_vocab_size, + org_vocab_size=args.vocab_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=getattr(config, "num_query_groups", config.num_attention_heads), + head_dim=getattr(config, "kv_channels", config.hidden_size // config.num_attention_heads), + rms_norm_eps=getattr(config, "layernorm_epsilon", 1e-6), + rotary_base=getattr(config, "rotary_base", 10000.0), + intermediate_size=args.dspark_intermediate_size, + ) + policy_embed = getattr(model.embedding, "word_embeddings", model.embedding) + policy_lm_head = getattr(model, "output_layer", None) + if policy_lm_head is None or getattr(policy_lm_head, "weight", None) is None: + policy_lm_head = policy_embed + model.draft_model = build_dspark_model( + dspark_config, + policy_embed, + policy_lm_head, + args.dspark_pretrained_model, + ) + for param in model.draft_model.parameters(): + param.grad_norm_group = "dspark" diff --git a/vime/backends/megatron_utils/fp8_helpers.py b/vime/backends/megatron_utils/fp8_helpers.py deleted file mode 100644 index 23980615e..000000000 --- a/vime/backends/megatron_utils/fp8_helpers.py +++ /dev/null @@ -1,72 +0,0 @@ -"""FP8 / UE8M0 quantization helpers for megatron → vLLM weight transfer. - -All symbols fall back to ``None`` when vLLM's deep_gemm helpers are not -available, which disables the UE8M0 requantization path. -""" - -import torch - -try: - import vllm.third_party.deep_gemm.utils.layout as _deep_gemm_layout - from vllm.utils.deep_gemm import get_tma_aligned_size as _get_tma_aligned_size - - _HAS_DEEP_GEMM = True -except ImportError: - _deep_gemm_layout = None - _get_tma_aligned_size = None - _HAS_DEEP_GEMM = False - -try: - from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used as _vllm_is_e8m0 - from vllm.utils.deep_gemm import per_block_cast_to_fp8 as _vllm_per_block_cast -except ImportError: - _vllm_is_e8m0 = lambda: False # noqa: E731 - _vllm_per_block_cast = None - - -def should_deepgemm_weight_requant_ue8m0(weight_block_size) -> bool: - return weight_block_size is not None and _vllm_is_e8m0() - - -def quant_weight_ue8m0( - weight_dequant: torch.Tensor, - weight_block_size: list[int], -): - assert weight_block_size == [128, 128] - assert weight_dequant.dtype == torch.bfloat16, f"{weight_dequant.dtype=} {weight_dequant.shape=}" - *batch_dims, n, k = weight_dequant.shape - flat = weight_dequant.view(-1, k) - out_w_flat, out_s_flat = _vllm_per_block_cast(flat, block_size=[128, 128], use_ue8m0=True) - out_w = out_w_flat.view(*batch_dims, n, k) - from math import ceil - - out_s = out_s_flat.view( - *batch_dims, - ceil(n / weight_block_size[0]), - ceil(k / weight_block_size[1]), - ) - return out_w, out_s - - -def transform_scale_ue8m0(sf: torch.Tensor, mn: int, use_torch_impl: bool = False): - if _deep_gemm_layout is None: - raise RuntimeError("deep_gemm not installed; UE8M0 requantization unavailable.") - get_fn = _deep_gemm_layout.get_mn_major_tma_aligned_packed_ue8m0_tensor - sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128) - sf = get_fn(sf) - if sf.shape[-1] == 1: - get_tma_aligned_size = _get_tma_aligned_size # pre-imported with fallback - - aligned_mn = get_tma_aligned_size(sf.shape[-2], sf.element_size()) - if sf.stride(-1) != aligned_mn: - new_stride = list(sf.stride()) - new_stride[-1] = aligned_mn - sf = sf.as_strided(sf.shape, tuple(new_stride)) - return sf - - -__all__ = [ - "quant_weight_ue8m0", - "transform_scale_ue8m0", - "should_deepgemm_weight_requant_ue8m0", -] diff --git a/vime/backends/megatron_utils/hf_checkpoint_saver.py b/vime/backends/megatron_utils/hf_checkpoint_saver.py index 76f0a6ef6..c442706ac 100644 --- a/vime/backends/megatron_utils/hf_checkpoint_saver.py +++ b/vime/backends/megatron_utils/hf_checkpoint_saver.py @@ -1,5 +1,6 @@ import json import logging +import math import os import shutil from pathlib import Path @@ -7,6 +8,8 @@ import torch +from vime.utils import accelerator + logger = logging.getLogger(__name__) _HF_WEIGHT_FILE_NAMES = { @@ -18,27 +21,38 @@ _HF_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".ckpt", ".msgpack") -def save_hf_model_direct(args, rollout_id: int, model) -> None: - """Save a Megatron model as an HF safetensors checkpoint without Megatron Bridge.""" +def save_hf_model_to_path( + args, + output_dir: str | Path, + model, + *, + model_name: str | None = None, + quantization_config: dict[str, Any] | None = None, + progress_desc: str = "Save HF checkpoint", +) -> None: + """Save a Megatron model as an HF safetensors checkpoint.""" + path = Path(output_dir) + hf_checkpoint = Path(args.hf_checkpoint).resolve() + save_path = path.resolve() + if hf_checkpoint == save_path: + raise ValueError("HF save output path must not point to the same directory as --hf-checkpoint") + if not hf_checkpoint.is_dir(): + raise ValueError( + f"--hf-checkpoint must be a local directory when saving raw HuggingFace weights: {args.hf_checkpoint}" + ) + import torch.distributed as dist from transformers import AutoConfig from .update_weight.common import named_params_and_buffers from .update_weight.hf_weight_iterator_direct import HfWeightIteratorDirect - path = Path(args.save_hf.format(rollout_id=rollout_id)) is_save_rank = _is_global_rank_zero() - hf_checkpoint = Path(args.hf_checkpoint).resolve() - save_path = path.resolve() - if hf_checkpoint == save_path: - raise ValueError("--save-hf must not point to the same directory as --hf-checkpoint") - if not hf_checkpoint.is_dir(): - raise ValueError(f"--hf-checkpoint must be a local directory when using raw --save-hf: {args.hf_checkpoint}") setup_error = None if is_save_rank: try: - logger.info("Saving model in HuggingFace format to %s with raw Megatron-to-HF conversion", path) + logger.info("Saving model in HuggingFace format to %s", path) path.mkdir(parents=True, exist_ok=True) _clear_existing_hf_weights(path) _copy_hf_assets(args.hf_checkpoint, path) @@ -49,18 +63,21 @@ def save_hf_model_direct(args, rollout_id: int, model) -> None: metadata_error = None payload: list[Any] = [None] - if is_save_rank: - try: - hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) - payload = [ - ( - type(hf_config).__name__.lower() if args.model_name is None else args.model_name, - getattr(hf_config, "quantization_config", None), - ) - ] - except Exception as e: - metadata_error = repr(e) - _raise_if_rank_zero_failed("load HuggingFace conversion metadata", metadata_error) + if model_name is not None: + payload = [(model_name, quantization_config)] + else: + if is_save_rank: + try: + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + payload = [ + ( + type(hf_config).__name__.lower() if args.model_name is None else args.model_name, + getattr(hf_config, "quantization_config", None), + ) + ] + except Exception as e: + metadata_error = repr(e) + _raise_if_rank_zero_failed("load HuggingFace conversion metadata", metadata_error) if dist.is_available() and dist.is_initialized(): dist.broadcast_object_list(payload, src=0) @@ -71,30 +88,44 @@ def save_hf_model_direct(args, rollout_id: int, model) -> None: model=model, model_name=model_name, quantization_config=quantization_config, + transform_ue8m0=False, ) - megatron_local_weights = dict(named_params_and_buffers(args, model, convert_to_global_name=True)) - writer = _SafetensorShardWriter(path, enabled=is_save_rank) - - for hf_named_tensors in hf_weight_iterator.get_hf_weight_chunks( - megatron_local_weights, progress_desc="Save HF checkpoint" + megatron_local_weights = dict(named_params_and_buffers(args, model)) + num_save_nodes, save_node_rank, is_writer_rank, writer_ranks = _get_node_save_layout(args) + if is_save_rank: + logger.info( + "Raw HuggingFace save will write shards from %d node writer rank(s): %s", + num_save_nodes, + writer_ranks, + ) + + writer = _SafetensorShardWriter(path, enabled=is_writer_rank) + pending_write = None + + for chunk_idx, hf_named_tensors in enumerate( + hf_weight_iterator.get_hf_weight_chunks( + megatron_local_weights, + progress_desc=progress_desc, + # Megatron-to-HF conversion is stateful for some parameters. For + # example, q_a_proj and kv_a_proj can land in adjacent chunks but + # must be emitted together for VLLM compatibility. Every node + # writer therefore has to observe every chunk so that pairs can + # cross chunk boundaries. Writers still only persist their + # modulo-assigned shards below; non-writer ranks skip conversion. + should_convert_chunk=lambda _idx: is_writer_rank, + ) ): - write_error = None - try: - writer.write(hf_named_tensors) - except Exception as e: - write_error = repr(e) - _raise_if_rank_zero_failed("write raw HuggingFace weight shard", write_error) - del hf_named_tensors - if torch.cuda.is_available(): - torch.cuda.ipc_collect() + if is_writer_rank and chunk_idx % num_save_nodes == save_node_rank: + pending_write = (chunk_idx, hf_named_tensors) + hf_named_tensors = None + else: + del hf_named_tensors - finalize_error = None - if is_save_rank: - try: - writer.finalize() - except Exception as e: - finalize_error = repr(e) - _raise_if_rank_zero_failed("finalize raw HuggingFace checkpoint", finalize_error) + if (chunk_idx + 1) % num_save_nodes == 0: + pending_write = _write_pending_chunk(writer, pending_write) + + pending_write = _write_pending_chunk(writer, pending_write) + _finalize_distributed_shards(path, writer.state()) if is_save_rank: logger.info("Successfully saved HuggingFace model to %s", path) @@ -108,28 +139,43 @@ def __init__(self, path: Path, *, enabled: bool) -> None: self.weight_map: dict[str, str] = {} self.shard_files: list[str] = [] - def write(self, named_tensors) -> None: + def write(self, named_tensors, shard_idx: int) -> None: if not self.enabled: return + assert shard_idx is not None, "shard_idx must be set when writing HF shards" from safetensors.torch import save_file state_dict = {} + total_size = 0 for name, tensor in named_tensors: if name in self.weight_map or name in state_dict: raise ValueError(f"Duplicate HF tensor while saving: {name}") - self.total_size += tensor.numel() * tensor.element_size() + total_size += tensor.numel() * tensor.element_size() state_dict[name] = _tensor_for_safetensors(tensor) if not state_dict: return - filename = f"model-{len(self.shard_files) + 1:05d}.safetensors" + filename = self._next_filename(shard_idx) + if (self.path / filename).exists(): + raise ValueError(f"Duplicate HF shard file while saving: {filename}") + save_file(state_dict, self.path / filename, metadata={"format": "pt"}) self.shard_files.append(filename) + self.total_size += total_size for name in state_dict: self.weight_map[name] = filename + def state(self) -> dict[str, Any]: + if not self.enabled: + return {"total_size": 0, "weight_map": {}, "shard_files": []} + return { + "total_size": self.total_size, + "weight_map": dict(self.weight_map), + "shard_files": list(self.shard_files), + } + def finalize(self) -> None: if not self.enabled: return @@ -148,6 +194,112 @@ def finalize(self) -> None: with open(self.path / "model.safetensors.index.json", "w", encoding="utf-8") as f: json.dump(index_data, f, indent=2) + def _next_filename(self, shard_idx: int) -> str: + assert shard_idx is not None, "shard_idx must be set when naming HF shards" + return f"model-{shard_idx + 1:05d}.safetensors" + + +def _write_pending_chunk( + writer: _SafetensorShardWriter, pending_write: tuple[int, Any] | None +) -> tuple[int, Any] | None: + if pending_write is not None: + shard_idx, named_tensors = pending_write + writer.write(named_tensors, shard_idx=shard_idx) + selected_accelerator = accelerator.initialize_accelerator() + if selected_accelerator is not None: + selected_accelerator.ipc_collect() + selected_accelerator.empty_cache() + + return None + + +def _finalize_distributed_shards(path: Path, local_state: dict[str, Any]) -> None: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + states = [None] * dist.get_world_size() + dist.all_gather_object(states, local_state) + else: + states = [local_state] + + _finalize_local_shards(path, local_state, states, write_index=_is_global_rank_zero()) + + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +def _finalize_local_shards( + path: Path, + local_state: dict[str, Any], + shard_states: list[dict[str, Any] | None], + *, + write_index: bool, +) -> None: + """Rename this rank's shard files per the global plan; optionally write the index. + + The plan is deterministic from the gathered states, so each rank renames only + its own files: on a non-POSIX shared filesystem another rank's unpublished + writes are not visible, let alone renamable. + """ + rename_map, index_data = _plan_shard_finalization(shard_states) + for old_name in local_state.get("shard_files", []): + os.replace(path / old_name, path / rename_map[old_name]) + if write_index: + with open(path / "model.safetensors.index.json", "w", encoding="utf-8") as f: + json.dump(index_data, f, indent=2) + + +def _plan_shard_finalization( + shard_states: list[dict[str, Any] | None], +) -> tuple[dict[str, str], dict[str, Any]]: + """Compute the shard rename map and index from every rank's gathered state.""" + shard_files = [] + total_size = 0 + raw_weight_map = {} + + for state in shard_states: + if not state: + continue + + total_size += state.get("total_size", 0) + for filename in state.get("shard_files", []): + if filename in shard_files: + raise ValueError(f"Duplicate HF shard file while finalizing: {filename}") + shard_files.append(filename) + + for name, filename in state.get("weight_map", {}).items(): + if name in raw_weight_map: + raise ValueError(f"Duplicate HF tensor while finalizing: {name}") + raw_weight_map[name] = filename + + if not shard_files: + raise ValueError("No HF tensors were produced while saving") + + shard_files = sorted(shard_files, key=_shard_filename_sort_key) + total_files = len(shard_files) + rename_map = {} + for idx, old_name in enumerate(shard_files, start=1): + rename_map[old_name] = f"model-{idx:05d}-of-{total_files:05d}.safetensors" + + final_weight_map = {} + for name, filename in raw_weight_map.items(): + if filename not in rename_map: + raise ValueError(f"HF tensor {name} points to missing shard file {filename}") + final_weight_map[name] = rename_map[filename] + + index_data = {"metadata": {"total_size": total_size}, "weight_map": final_weight_map} + return rename_map, index_data + + +def _shard_filename_sort_key(filename: str) -> tuple[float, str]: + prefix = "model-" + suffix = ".safetensors" + if filename.startswith(prefix) and filename.endswith(suffix): + middle = filename[len(prefix) : -len(suffix)] + if middle.isdigit(): + return int(middle), filename + return math.inf, filename + def _tensor_for_safetensors(tensor: torch.Tensor) -> torch.Tensor: tensor = tensor.detach() @@ -187,6 +339,24 @@ def _is_global_rank_zero() -> bool: return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 +def _get_node_save_layout(args) -> tuple[int, int, bool, list[int]]: + import torch.distributed as dist + + if not (dist.is_available() and dist.is_initialized()): + return 1, 0, True, [0] + + world_size = dist.get_world_size() + rank = dist.get_rank() + gpus_per_node = int(getattr(args, "actor_num_gpus_per_node", None) or getattr(args, "num_gpus_per_node", 1) or 1) + gpus_per_node = max(1, gpus_per_node) + inferred_nodes = max(1, math.ceil(world_size / gpus_per_node)) + configured_nodes = int(getattr(args, "actor_num_nodes", None) or inferred_nodes) + num_nodes = max(1, min(configured_nodes, inferred_nodes)) + writer_ranks = [node * gpus_per_node for node in range(num_nodes) if node * gpus_per_node < world_size] + node_rank = min(rank // gpus_per_node, num_nodes - 1) + return len(writer_ranks), node_rank, rank in writer_ranks, writer_ranks + + def _raise_if_rank_zero_failed(context: str, error: str | None) -> None: import torch.distributed as dist diff --git a/vime/backends/megatron_utils/hf_to_megatron/__init__.py b/vime/backends/megatron_utils/hf_to_megatron/__init__.py new file mode 100644 index 000000000..a5398ae97 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/__init__.py @@ -0,0 +1,48 @@ +from pathlib import Path + +from transformers import AutoConfig + +from .common import load_model_hf_weights +from .deepseek import deepseek_hf_tensor +from .glm import glm4_hf_tensor, glm4_moe_hf_tensor +from .qwen import mimo_hf_tensor, minimax_m2_hf_tensor, qwen_hf_tensor, qwen_moe_hf_tensor +from .qwen3_5 import qwen3_5_hf_tensor +from .qwen3_next import qwen3_next_hf_tensor +from .qwen3_omni import qwen3_omni_hf_tensor +from .qwen3_vl import qwen3_vl_hf_tensor + +_LOADERS = { + "deepseek_v3": deepseek_hf_tensor, + "deepseek_v32": deepseek_hf_tensor, + "glm4": glm4_hf_tensor, + "glm4_moe": glm4_moe_hf_tensor, + "glm4_moe_lite": deepseek_hf_tensor, + "glm_moe_dsa": deepseek_hf_tensor, + "kimi_k2": deepseek_hf_tensor, + "llama": qwen_hf_tensor, + "mimo": mimo_hf_tensor, + "minimax_m2": minimax_m2_hf_tensor, + "qwen2": qwen_hf_tensor, + "qwen2_moe": qwen_moe_hf_tensor, + "qwen3": qwen_hf_tensor, + "qwen3_5": qwen3_5_hf_tensor, + "qwen3_5_moe": qwen3_5_hf_tensor, + "qwen3_moe": qwen_moe_hf_tensor, + "qwen3_next": qwen3_next_hf_tensor, + "qwen3_omni_moe": qwen3_omni_hf_tensor, + "qwen3_vl": qwen3_vl_hf_tensor, +} + + +def supports_hf_weight_loading(path: str | Path) -> bool: + config = AutoConfig.from_pretrained(path, trust_remote_code=True) + return config.model_type in _LOADERS + + +def load_hf_weights(args, model, path: str | Path) -> None: + config = AutoConfig.from_pretrained(path, trust_remote_code=True) + try: + get_hf_tensor = _LOADERS[config.model_type] + except KeyError as exc: + raise ValueError(f"Unsupported HuggingFace model type: {config.model_type}") from exc + load_model_hf_weights(args, model, path, config, get_hf_tensor) diff --git a/vime/backends/megatron_utils/hf_to_megatron/common.py b/vime/backends/megatron_utils/hf_to_megatron/common.py new file mode 100644 index 000000000..8adff4198 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/common.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import functools +import json +from collections.abc import Callable +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from vime.platforms import current_platform + + +class SafetensorReader: + def __init__(self, path: str | Path): + self.path = Path(path) + index_path = self.path / "model.safetensors.index.json" + if index_path.is_file(): + with index_path.open() as index_file: + self.weight_map = json.load(index_file)["weight_map"] + else: + files = sorted(self.path.glob("*.safetensors")) + if not files: + raise FileNotFoundError(f"No safetensors checkpoint found in {self.path}") + self.weight_map = {} + for file in files: + with safe_open(file, framework="pt", device="cpu") as tensors: + self.weight_map.update(dict.fromkeys(tensors.keys(), file.name)) + self._files = {} + + def __contains__(self, name: str) -> bool: + return name in self.weight_map + + @functools.lru_cache(maxsize=1) # noqa: B019 - cache belongs to this reader instance + def get_tensor(self, name: str) -> torch.Tensor: + try: + filename = self.weight_map[name] + except KeyError as exc: + raise KeyError(f"HuggingFace checkpoint does not contain {name!r}") from exc + if filename not in self._files: + self._files[filename] = safe_open(self.path / filename, framework="pt", device="cpu") + tensor = self._files[filename].get_tensor(name) + scale_name = f"{name}_scale_inv" + if tensor.element_size() == 1 and scale_name in self: + scale_file = self.weight_map[scale_name] + if scale_file not in self._files: + self._files[scale_file] = safe_open(self.path / scale_file, framework="pt", device="cpu") + scale = self._files[scale_file].get_tensor(scale_name).to(torch.bfloat16) + rows, columns = tensor.shape + block_rows, block_columns = scale.shape + tensor = F.pad( + tensor.to(torch.bfloat16), + (0, block_columns * 128 - columns, 0, block_rows * 128 - rows), + ) + tensor = tensor.view(block_rows, 128, block_columns, 128) + tensor.mul_(scale[:, None, :, None]) + tensor = tensor.reshape(block_rows * 128, block_columns * 128)[:rows, :columns] + return tensor + + +def strip_mcore_wrappers(name: str) -> str: + while name.startswith("module."): + name = name.removeprefix("module.") + return name.removeprefix("language_model.") + + +def text_config(config): + return getattr(config, "text_config", config) + + +def merge_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, config) -> torch.Tensor: + config = text_config(config) + num_groups = config.num_key_value_heads + num_heads = config.num_attention_heads + head_dim = getattr(config, "head_dim", None) or config.hidden_size // num_heads + trailing_shape = q.shape[1:] + q = q.reshape(num_groups, num_heads // num_groups * head_dim, *trailing_shape) + k = k.reshape(num_groups, head_dim, *trailing_shape) + v = v.reshape(num_groups, head_dim, *trailing_shape) + return torch.cat((q, k, v), dim=1).reshape(-1, *trailing_shape).contiguous() + + +def merge_gate_up(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat((gate, up), dim=0) + + +def _tensor_parallel_shard( + name: str, + tensor: torch.Tensor, + *, + parallel_size: int, + parallel_rank: int, + partition_dim: int, + partition_stride: int, +) -> torch.Tensor: + if parallel_size == 1: + return tensor + + if "linear_fc1.weight" in name or "linear_fc1.bias" in name: + gate, up = tensor.chunk(2, dim=partition_dim) + gate = torch.chunk(gate, parallel_size, dim=partition_dim)[parallel_rank] + up = torch.chunk(up, parallel_size, dim=partition_dim)[parallel_rank] + return torch.cat((gate, up), dim=partition_dim).contiguous() + + if "linear_fc2.weight" in name and partition_dim == 0: + partition_dim = 1 + + chunks = torch.chunk(tensor, parallel_size * partition_stride, dim=partition_dim) + return torch.cat(chunks[parallel_rank::parallel_size], dim=partition_dim).contiguous() + + +def shard_mcore_tensor(name: str, tensor: torch.Tensor, parameter: torch.Tensor) -> torch.Tensor: + from megatron.core import mpu + + if ( + not getattr(parameter, "tensor_model_parallel", False) + or getattr(parameter, "parallel_mode", None) == "duplicated" + ): + return tensor + + if ".experts." in name: + parallel_size = mpu.get_expert_tensor_parallel_world_size() + parallel_rank = mpu.get_expert_tensor_parallel_rank() + else: + parallel_size = mpu.get_tensor_model_parallel_world_size() + parallel_rank = mpu.get_tensor_model_parallel_rank() + + return _tensor_parallel_shard( + name, + tensor, + parallel_size=parallel_size, + parallel_rank=parallel_rank, + partition_dim=current_platform().megatron.adjust_tp_partition_dim(name, parameter.partition_dim), + partition_stride=parameter.partition_stride, + ) + + +def _pad_vocab(args, name: str, tensor: torch.Tensor) -> torch.Tensor: + if not (name.endswith("embedding.word_embeddings.weight") or name.endswith("output_layer.weight")): + return tensor + padded_size = getattr(args, "padded_vocab_size", None) + if padded_size is None or tensor.shape[0] >= padded_size: + return tensor + return F.pad(tensor, (0, 0, 0, padded_size - tensor.shape[0])) + + +def load_model_hf_weights( + args, + model, + path: str | Path, + config, + get_hf_tensor: Callable[[str, SafetensorReader, object], torch.Tensor], +) -> None: + from vime.backends.megatron_utils.update_weight.common import named_params_and_buffers + + reader = SafetensorReader(path) + with torch.no_grad(): + for name, parameter in named_params_and_buffers(args, model): + tensor = get_hf_tensor(name, reader, config) + if name.endswith("output_layer.weight") and parameter.shape[0] == 1 and tensor.shape[0] != 1: + continue + tensor = shard_mcore_tensor(name, _pad_vocab(args, name, tensor), parameter) + if tensor.shape != parameter.shape: + raise ValueError( + f"Shape mismatch loading {name}: HuggingFace {tuple(tensor.shape)}, " + f"Megatron {tuple(parameter.shape)}" + ) + parameter.copy_(tensor.to(device=parameter.device, dtype=parameter.dtype)) diff --git a/vime/backends/megatron_utils/hf_to_megatron/deepseek.py b/vime/backends/megatron_utils/hf_to_megatron/deepseek.py new file mode 100644 index 000000000..e7f8fbe37 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/deepseek.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import re + +import torch + +from .common import SafetensorReader, merge_gate_up, strip_mcore_wrappers + + +def _direct(name: str, reader: SafetensorReader, config) -> torch.Tensor | None: + mapping = { + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + "output_layer.weight": ( + "model.embed_tokens.weight" + if getattr(config, "tie_word_embeddings", False) or "lm_head.weight" not in reader + else "lm_head.weight" + ), + } + return reader.get_tensor(mapping[name]) if name in mapping else None + + +def _dsa_reorder(name: str, tensor: torch.Tensor) -> torch.Tensor: + if name == "self_attention.wq_b.weight": + tensor = tensor.view(-1, 128, tensor.shape[-1]) + return torch.cat((tensor[:, 64:], tensor[:, :64]), dim=1).flatten(0, 1) + if name in { + "self_attention.wk.weight", + "self_attention.k_norm.weight", + "self_attention.k_norm.bias", + }: + return torch.cat((tensor[64:], tensor[:64]), dim=0) + return tensor + + +def _layer_tensor( + layer: int, + rest: str, + reader: SafetensorReader, + *, + dsa: bool, +) -> torch.Tensor: + prefix = f"model.layers.{layer}" + mapping = { + "input_layernorm.weight": "input_layernorm.weight", + "self_attention.linear_qkv.layer_norm_weight": "input_layernorm.weight", + "self_attention.linear_proj.weight": "self_attn.o_proj.weight", + "self_attention.linear_q_proj.weight": "self_attn.q_proj.weight", + "self_attention.linear_kv_down_proj.weight": "self_attn.kv_a_proj_with_mqa.weight", + "self_attention.linear_kv_up_proj.layer_norm_weight": "self_attn.kv_a_layernorm.weight", + "self_attention.linear_kv_up_proj.weight": "self_attn.kv_b_proj.weight", + "self_attention.linear_q_down_proj.weight": "self_attn.q_a_proj.weight", + "self_attention.linear_q_up_proj.weight": "self_attn.q_b_proj.weight", + "self_attention.linear_q_up_proj.layer_norm_weight": "self_attn.q_a_layernorm.weight", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "pre_mlp_layernorm.weight": "post_attention_layernorm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + "mlp.shared_experts.linear_fc2.weight": "mlp.shared_experts.down_proj.weight", + "mlp.router.weight": "mlp.gate.weight", + "mlp.router.expert_bias": "mlp.gate.e_score_correction_bias", + } + if dsa: + mapping |= { + "self_attention.wq_b.weight": "self_attn.indexer.wq_b.weight", + "self_attention.wk.weight": "self_attn.indexer.wk.weight", + "self_attention.weights_proj.weight": "self_attn.indexer.weights_proj.weight", + "self_attention.k_norm.weight": "self_attn.indexer.k_norm.weight", + "self_attention.k_norm.bias": "self_attn.indexer.k_norm.bias", + } + if rest in mapping: + return ( + _dsa_reorder(rest, reader.get_tensor(f"{prefix}.{mapping[rest]}")) + if dsa + else reader.get_tensor(f"{prefix}.{mapping[rest]}") + ) + if rest == "mlp.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.up_proj.weight"), + ) + if rest == "mlp.shared_experts.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.shared_experts.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.shared_experts.up_proj.weight"), + ) + match = re.fullmatch(r"mlp\.experts\.linear_fc([12])\.weight(\d+)", rest) + if match: + projection, expert = match.groups() + if projection == "1": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.experts.{expert}.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.experts.{expert}.up_proj.weight"), + ) + return reader.get_tensor(f"{prefix}.mlp.experts.{expert}.down_proj.weight") + raise KeyError(f"Unsupported DeepSeek Megatron layer parameter {rest!r}") + + +def deepseek_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct(name, reader, config)) is not None: + return tensor + + mtp = re.fullmatch(r"mtp\.layers\.(\d+)\.(.+)", name) + if mtp: + mtp_layer, rest = mtp.groups() + layer = config.num_hidden_layers + int(mtp_layer) + mapping = { + "eh_proj.weight": "eh_proj.weight", + "enorm.weight": "enorm.weight", + "hnorm.weight": "hnorm.weight", + "final_layernorm.weight": "shared_head.norm.weight", + } + if rest in mapping: + return reader.get_tensor(f"model.layers.{layer}.{mapping[rest]}") + rest = rest.removeprefix("transformer_layer.") + else: + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported DeepSeek Megatron parameter {name!r}") + layer, rest = int(match.group(1)), match.group(2) + + return _layer_tensor( + int(layer), + rest, + reader, + dsa=config.model_type in {"deepseek_v32", "glm_moe_dsa"}, + ) diff --git a/vime/backends/megatron_utils/hf_to_megatron/glm.py b/vime/backends/megatron_utils/hf_to_megatron/glm.py new file mode 100644 index 000000000..a7e973394 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/glm.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import re + +import torch + +from .common import SafetensorReader, merge_gate_up, strip_mcore_wrappers +from .qwen import _attention_tensor, _direct_tensor, _qwen_moe_layer_tensor + + +def glm4_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported GLM-4 Megatron parameter {name!r}") + layer, rest = match.groups() + prefix = f"model.layers.{layer}" + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + mapping = { + "mlp.linear_fc1.weight": "mlp.gate_up_proj.weight", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + "post_self_attn_layernorm.weight": "post_self_attn_layernorm.weight", + "post_mlp_layernorm.weight": "post_mlp_layernorm.weight", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + raise KeyError(f"Unsupported GLM-4 Megatron parameter {name!r}") + + +def _glm4_moe_layer_tensor(rest: str, prefix: str, reader: SafetensorReader, config) -> torch.Tensor: + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + if rest == "mlp.shared_experts.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.shared_experts.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.shared_experts.up_proj.weight"), + ) + if rest == "mlp.shared_experts.linear_fc2.weight": + return reader.get_tensor(f"{prefix}.mlp.shared_experts.down_proj.weight") + if (tensor := _qwen_moe_layer_tensor(rest, prefix, reader)) is not None: + return tensor + mapping = { + "mlp.router.expert_bias": "mlp.gate.e_score_correction_bias", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + if rest == "mlp.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.up_proj.weight"), + ) + raise KeyError(f"Unsupported GLM-4 MoE Megatron layer parameter {rest!r}") + + +def glm4_moe_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + mtp = re.fullmatch(r"mtp\.layers\.(\d+)\.(.+)", name) + if mtp: + mtp_layer, rest = mtp.groups() + layer = config.num_hidden_layers + int(mtp_layer) + mapping = { + "enorm.weight": "enorm.weight", + "hnorm.weight": "hnorm.weight", + "eh_proj.weight": "eh_proj.weight", + "final_layernorm.weight": "shared_head.norm.weight", + } + if rest in mapping: + return reader.get_tensor(f"model.layers.{layer}.{mapping[rest]}") + return _glm4_moe_layer_tensor( + rest.removeprefix("transformer_layer."), + f"model.layers.{layer}", + reader, + config, + ) + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported GLM-4 MoE Megatron parameter {name!r}") + layer, rest = match.groups() + return _glm4_moe_layer_tensor(rest, f"model.layers.{layer}", reader, config) diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen.py b/vime/backends/megatron_utils/hf_to_megatron/qwen.py new file mode 100644 index 000000000..29cebc760 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/qwen.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import re + +import torch + +from .common import SafetensorReader, merge_gate_up, merge_qkv, strip_mcore_wrappers + + +def _direct_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor | None: + mapping = { + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + "output_layer.weight": ( + "model.embed_tokens.weight" + if getattr(config, "tie_word_embeddings", False) or "lm_head.weight" not in reader + else "lm_head.weight" + ), + } + return reader.get_tensor(mapping[name]) if name in mapping else None + + +def _layer(name: str) -> tuple[int, str]: + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported Megatron parameter {name!r}") + return int(match.group(1)), match.group(2) + + +def _attention_tensor( + rest: str, + prefix: str, + reader: SafetensorReader, + config, +) -> torch.Tensor | None: + mapping = { + "self_attention.linear_proj.weight": "self_attn.o_proj.weight", + "self_attention.linear_proj.bias": "self_attn.o_proj.bias", + "self_attention.linear_qkv.layer_norm_weight": "input_layernorm.weight", + "self_attention.q_layernorm.weight": "self_attn.q_norm.weight", + "self_attention.k_layernorm.weight": "self_attn.k_norm.weight", + "self_attention.core_attention.softmax_offset": "self_attn.sinks", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + match = re.fullmatch(r"self_attention\.linear_qkv\.(weight|bias)", rest) + if match: + suffix = match.group(1) + return merge_qkv( + *(reader.get_tensor(f"{prefix}.self_attn.{projection}_proj.{suffix}") for projection in "qkv"), + config, + ) + return None + + +def qwen_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + + layer, rest = _layer(name) + prefix = f"model.layers.{layer}" + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + + mapping = { + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "pre_mlp_layernorm.weight": "post_attention_layernorm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + if rest == "mlp.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.up_proj.weight"), + ) + raise KeyError(f"Unsupported Qwen/Llama Megatron parameter {name!r}") + + +def _qwen_moe_layer_tensor( + rest: str, + prefix: str, + reader: SafetensorReader, +) -> torch.Tensor | None: + mapping = { + "pre_mlp_layernorm.weight": "post_attention_layernorm.weight", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "mlp.router.weight": "mlp.gate.weight", + "mlp.router.expert_bias": "mlp.gate.e_score_correction_bias", + "mlp.shared_experts.linear_fc2.weight": "mlp.shared_expert.down_proj.weight", + "mlp.shared_experts.gate_weight": "mlp.shared_expert_gate.weight", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + if rest in {"mlp.shared_experts.linear_fc1.weight", "shared_experts.linear_fc1.weight"}: + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.shared_expert.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.shared_expert.up_proj.weight"), + ) + if rest == "shared_experts.linear_fc2.weight": + return reader.get_tensor(f"{prefix}.mlp.shared_expert.down_proj.weight") + if rest == "shared_experts.gate_weight": + return reader.get_tensor(f"{prefix}.mlp.shared_expert_gate.weight") + + match = re.fullmatch(r"mlp\.experts\.linear_fc([12])\.(weight|bias)(\d+)", rest) + if match: + projection, kind, expert = match.groups() + if projection == "1": + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.experts.{expert}.gate_proj.{kind}"), + reader.get_tensor(f"{prefix}.mlp.experts.{expert}.up_proj.{kind}"), + ) + return reader.get_tensor(f"{prefix}.mlp.experts.{expert}.down_proj.{kind}") + return None + + +def qwen_moe_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + + layer, rest = _layer(name) + prefix = f"model.layers.{layer}" + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + if (tensor := _qwen_moe_layer_tensor(rest, prefix, reader)) is not None: + return tensor + return qwen_hf_tensor(name, reader, config) + + +def mimo_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + match = re.fullmatch(r"mtp\.layers\.(\d+)\.(.+)", name) + if not match: + return qwen_hf_tensor(name, reader, config) + + layer, rest = match.groups() + mapping = { + "enorm.weight": "token_layernorm.weight", + "hnorm.weight": "hidden_layernorm.weight", + "eh_proj.weight": "input_proj.weight", + "final_layernorm.weight": "final_layernorm.weight", + } + if rest in mapping: + tensor = reader.get_tensor(f"model.mtp_layers.{layer}.{mapping[rest]}") + if rest == "eh_proj.weight": + tensor = torch.cat(tensor.chunk(2, dim=1)[::-1], dim=1) + return tensor + + proxy = f"decoder.layers.{layer}." + rest.removeprefix("transformer_layer.") + hf_prefix = f"model.mtp_layers.{layer}" + _, proxy_rest = _layer(proxy) + if (tensor := _attention_tensor(proxy_rest, hf_prefix, reader, config)) is not None: + return tensor + mapping = { + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + } + if proxy_rest in mapping: + return reader.get_tensor(f"{hf_prefix}.{mapping[proxy_rest]}") + if proxy_rest == "mlp.linear_fc1.weight": + return merge_gate_up( + reader.get_tensor(f"{hf_prefix}.mlp.gate_proj.weight"), + reader.get_tensor(f"{hf_prefix}.mlp.up_proj.weight"), + ) + raise KeyError(f"Unsupported MiMo Megatron parameter {name!r}") + + +def minimax_m2_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + layer, rest = _layer(name) + prefix = f"model.layers.{layer}" + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + mapping = { + "pre_mlp_layernorm.weight": "post_attention_layernorm.weight", + "mlp.router.weight": "block_sparse_moe.gate.weight", + "mlp.router.expert_bias": "block_sparse_moe.e_score_correction_bias", + "self_attention.q_norm.weight": "self_attn.q_norm.weight", + "self_attention.k_norm.weight": "self_attn.k_norm.weight", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + match = re.fullmatch(r"mlp\.experts\.linear_fc([12])\.weight(\d+)", rest) + if match: + projection, expert = match.groups() + if projection == "1": + return merge_gate_up( + reader.get_tensor(f"{prefix}.block_sparse_moe.experts.{expert}.w1.weight"), + reader.get_tensor(f"{prefix}.block_sparse_moe.experts.{expert}.w3.weight"), + ) + return reader.get_tensor(f"{prefix}.block_sparse_moe.experts.{expert}.w2.weight") + raise KeyError(f"Unsupported MiniMax-M2 Megatron parameter {name!r}") diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen3_5.py b/vime/backends/megatron_utils/hf_to_megatron/qwen3_5.py new file mode 100644 index 000000000..254290a22 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/qwen3_5.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import re + +import torch + +from .common import SafetensorReader, strip_mcore_wrappers + + +def _merge_qkv(reader: SafetensorReader, prefix: str, text_config, suffix: str) -> torch.Tensor: + q = reader.get_tensor(f"{prefix}.q_proj.{suffix}") + k = reader.get_tensor(f"{prefix}.k_proj.{suffix}") + v = reader.get_tensor(f"{prefix}.v_proj.{suffix}") + num_groups = text_config.num_key_value_heads + queries_per_group = text_config.num_attention_heads // num_groups + head_dim = text_config.head_dim + + trailing_shape = q.shape[1:] + q = q.reshape(num_groups, queries_per_group, 2, head_dim, *trailing_shape).transpose(1, 2) + q = q.flatten(1, 3) + k = k.reshape(num_groups, head_dim, *trailing_shape) + v = v.reshape(num_groups, head_dim, *trailing_shape) + return torch.cat((q, k, v), dim=1).reshape(-1, *trailing_shape).contiguous() + + +def qwen3_5_hf_tensor(name: str, reader: SafetensorReader, hf_config) -> torch.Tensor: + """Return the full, unsharded MCore tensor for a Qwen3.5 parameter name.""" + + name = strip_mcore_wrappers(name) + if name.startswith("model.visual."): + return reader.get_tensor(name) + name = name.removeprefix("language_model.") + + text_config = getattr(hf_config, "text_config", hf_config) + direct_mapping = { + "embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight", + "decoder.final_layernorm.weight": "model.language_model.norm.weight", + "output_layer.weight": ( + "model.language_model.embed_tokens.weight" + if getattr(hf_config, "tie_word_embeddings", False) or getattr(text_config, "tie_word_embeddings", False) + else "lm_head.weight" + ), + } + if name in direct_mapping: + return reader.get_tensor(direct_mapping[name]) + + mtp_match = re.fullmatch(r"mtp\.layers\.(\d+)\.(.+)", name) + if mtp_match: + is_mtp = True + mtp_layer, rest = mtp_match.groups() + direct_mtp = { + "eh_proj.weight": "mtp.fc.weight", + "enorm.weight": "mtp.pre_fc_norm_embedding.weight", + "hnorm.weight": "mtp.pre_fc_norm_hidden.weight", + "final_layernorm.weight": "mtp.norm.weight", + } + if rest in direct_mtp: + return reader.get_tensor(direct_mtp[rest]) + rest = rest.removeprefix("transformer_layer.") + name = f"decoder.layers.{mtp_layer}." + rest + hf_layer_prefix = f"mtp.layers.{mtp_layer}" + else: + is_mtp = False + layer_match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not layer_match: + raise KeyError(f"Unsupported Qwen3.5 Megatron parameter {name!r}") + layer_idx, rest = layer_match.groups() + hf_layer_prefix = f"model.language_model.layers.{layer_idx}" + + if rest.startswith("self_attention.linear_attn."): + suffix = rest.removeprefix("self_attention.") + return reader.get_tensor(f"{hf_layer_prefix}.{suffix}") + if rest == "self_attention.input_layernorm.weight": + return reader.get_tensor(f"{hf_layer_prefix}.input_layernorm.weight") + if rest == "self_attention.linear_proj.weight": + return reader.get_tensor(f"{hf_layer_prefix}.self_attn.o_proj.weight") + if rest == "self_attention.linear_qkv.weight": + return _merge_qkv(reader, f"{hf_layer_prefix}.self_attn", text_config, "weight") + if rest == "self_attention.linear_qkv.bias": + return _merge_qkv(reader, f"{hf_layer_prefix}.self_attn", text_config, "bias") + if rest == "self_attention.linear_qkv.layer_norm_weight": + return reader.get_tensor(f"{hf_layer_prefix}.input_layernorm.weight") + if rest == "self_attention.q_layernorm.weight": + return reader.get_tensor(f"{hf_layer_prefix}.self_attn.q_norm.weight") + if rest == "self_attention.k_layernorm.weight": + return reader.get_tensor(f"{hf_layer_prefix}.self_attn.k_norm.weight") + + if rest in {"mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight"}: + return reader.get_tensor(f"{hf_layer_prefix}.post_attention_layernorm.weight") + if rest == "mlp.linear_fc1.weight": + gate = reader.get_tensor(f"{hf_layer_prefix}.mlp.gate_proj.weight") + up = reader.get_tensor(f"{hf_layer_prefix}.mlp.up_proj.weight") + return torch.cat((gate, up), dim=0) + if rest == "mlp.linear_fc2.weight": + return reader.get_tensor(f"{hf_layer_prefix}.mlp.down_proj.weight") + if rest == "mlp.router.weight": + return reader.get_tensor(f"{hf_layer_prefix}.mlp.gate.weight") + if rest == "mlp.router.expert_bias": + return reader.get_tensor(f"{hf_layer_prefix}.mlp.gate.e_score_correction_bias") + + expert_match = re.fullmatch(r"mlp\.experts\.linear_fc([12])(?:\.weight)?(\d+)?", rest) + if expert_match: + projection, expert_idx = expert_match.groups() + if is_mtp and expert_idx is not None: + prefix = f"{hf_layer_prefix}.mlp.experts.{expert_idx}" + if projection == "1": + return torch.cat( + ( + reader.get_tensor(f"{prefix}.gate_proj.weight"), + reader.get_tensor(f"{prefix}.up_proj.weight"), + ), + dim=0, + ) + return reader.get_tensor(f"{prefix}.down_proj.weight") + suffix = "gate_up_proj" if projection == "1" else "down_proj" + tensor = reader.get_tensor(f"{hf_layer_prefix}.mlp.experts.{suffix}") + return tensor if expert_idx is None else tensor[int(expert_idx)].contiguous() + + shared_mapping = { + "mlp.shared_experts.linear_fc1.weight": ("gate_proj.weight", "up_proj.weight"), + "mlp.shared_experts.linear_fc2.weight": ("down_proj.weight",), + "mlp.shared_experts.gate_weight": ("../shared_expert_gate.weight",), + } + if rest in shared_mapping: + tensors = [] + for suffix in shared_mapping[rest]: + if suffix.startswith("../"): + key = f"{hf_layer_prefix}.mlp.{suffix.removeprefix('../')}" + else: + key = f"{hf_layer_prefix}.mlp.shared_expert.{suffix}" + tensors.append(reader.get_tensor(key)) + return tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + + raise KeyError(f"Unsupported Qwen3.5 Megatron parameter {name!r}") diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen3_next.py b/vime/backends/megatron_utils/hf_to_megatron/qwen3_next.py new file mode 100644 index 000000000..f14d42227 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/qwen3_next.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import re + +import torch + +from .common import SafetensorReader, strip_mcore_wrappers +from .qwen import _attention_tensor, _direct_tensor, _qwen_moe_layer_tensor + +_DIRECT_ATTENTION = { + "input_layernorm.weight", + "linear_attn.A_log", + "linear_attn.conv1d.weight", + "linear_attn.dt_bias", + "linear_attn.in_proj_ba.weight", + "linear_attn.in_proj_qkvz.weight", + "linear_attn.norm.weight", + "linear_attn.out_proj.weight", + "self_attn.k_norm.weight", + "self_attn.k_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.q_proj.weight", + "self_attn.v_proj.weight", +} + + +def _qwen3_next_layer_tensor( + rest: str, + prefix: str, + reader: SafetensorReader, + config, +) -> torch.Tensor: + direct = rest.removeprefix("self_attention.") + if rest.startswith("self_attention.") and direct in _DIRECT_ATTENTION: + return reader.get_tensor(f"{prefix}.{direct}") + if rest in {"self_attention.linear_qkv.weight", "self_attention.linear_qkv.bias"}: + suffix = rest.rsplit(".", 1)[-1] + q, k, v = (reader.get_tensor(f"{prefix}.self_attn.{projection}_proj.{suffix}") for projection in "qkv") + text = getattr(config, "text_config", config) + groups = text.num_key_value_heads + queries_per_group = text.num_attention_heads // groups + head_dim = getattr(text, "head_dim", None) or text.hidden_size // text.num_attention_heads + trailing = q.shape[1:] + q = q.reshape(groups, queries_per_group, 2, head_dim, *trailing).transpose(1, 2).flatten(1, 3) + k = k.reshape(groups, head_dim, *trailing) + v = v.reshape(groups, head_dim, *trailing) + return torch.cat((q, k, v), dim=1).reshape(-1, *trailing).contiguous() + if (tensor := _attention_tensor(rest, prefix, reader, config)) is not None: + return tensor + if (tensor := _qwen_moe_layer_tensor(rest, prefix, reader)) is not None: + return tensor + raise KeyError(f"Unsupported Qwen3-Next Megatron layer parameter {rest!r}") + + +def qwen3_next_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + if (tensor := _direct_tensor(name, reader, config)) is not None: + return tensor + + mtp = re.fullmatch(r"mtp\.layers\.(\d+)\.(.+)", name) + if mtp: + layer, rest = mtp.groups() + mapping = { + "eh_proj.weight": "mtp.fc.weight", + "enorm.weight": "mtp.pre_fc_norm_embedding.weight", + "hnorm.weight": "mtp.pre_fc_norm_hidden.weight", + "final_layernorm.weight": "mtp.norm.weight", + } + if rest in mapping: + tensor = reader.get_tensor(mapping[rest]) + if rest == "eh_proj.weight": + tensor = torch.cat(tensor.chunk(2, dim=1)[::-1], dim=1) + return tensor + prefix = f"mtp.layers.{layer}" + rest = rest.removeprefix("transformer_layer.") + return _qwen3_next_layer_tensor(rest, prefix, reader, config) + + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported Qwen3-Next Megatron parameter {name!r}") + layer, rest = match.groups() + return _qwen3_next_layer_tensor(rest, f"model.layers.{layer}", reader, config) diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen3_omni.py b/vime/backends/megatron_utils/hf_to_megatron/qwen3_omni.py new file mode 100644 index 000000000..420243c58 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/qwen3_omni.py @@ -0,0 +1,21 @@ +from .common import SafetensorReader, strip_mcore_wrappers +from .qwen import qwen_moe_hf_tensor + + +class _ThinkerReader: + def __init__(self, reader: SafetensorReader) -> None: + self.reader = reader + + def __contains__(self, name: str) -> bool: + return f"thinker.{name}" in self.reader + + def get_tensor(self, name: str): + return self.reader.get_tensor(f"thinker.{name}") + + +def qwen3_omni_hf_tensor(name: str, reader: SafetensorReader, config): + name = strip_mcore_wrappers(name) + for model_prefix, hf_prefix in (("audio_model.", "audio_tower."), ("vision_model.", "visual.")): + if name.startswith(model_prefix): + return reader.get_tensor(f"thinker.{hf_prefix}{name.removeprefix(model_prefix)}") + return qwen_moe_hf_tensor(name, _ThinkerReader(reader), config.thinker_config.text_config) diff --git a/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py b/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py new file mode 100644 index 000000000..a5b350628 --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/qwen3_vl.py @@ -0,0 +1,26 @@ +from .common import strip_mcore_wrappers +from .qwen import qwen_hf_tensor + + +class _LanguageReader: + def __init__(self, reader): + self.reader = reader + + @staticmethod + def _key(name): + return name.replace("model.", "model.language_model.", 1) if name.startswith("model.") else name + + def __contains__(self, name): + return self._key(name) in self.reader + + def get_tensor(self, name): + return self.reader.get_tensor(self._key(name)) + + +def qwen3_vl_hf_tensor(name, reader, config): + name = strip_mcore_wrappers(name) + if name.startswith("model.visual."): + return reader.get_tensor(name) + if name == "output_layer.weight" and getattr(config, "tie_word_embeddings", False): + return reader.get_tensor("model.language_model.embed_tokens.weight") + return qwen_hf_tensor(name, _LanguageReader(reader), config.text_config) diff --git a/vime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu b/vime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu index a6e955490..f7df09987 100644 --- a/vime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu +++ b/vime/backends/megatron_utils/kernels/int4_qat/fake_int4_quant_cuda.cu @@ -1,7 +1,15 @@ #include #include +// HIP's __shfl_xor_sync requires a 64-bit mask (a wavefront is 64 lanes), so +// 0xFFFFFFFF does not compile. The reductions only ever span one 32-lane group, +// which the maskless __shfl_xor with width=32 expresses identically. +#if defined(__HIP_PLATFORM_AMD__) +#define WARP_XOR(val, mask) __shfl_xor((val), (mask), 32) +#else #define FINAL_MASK 0xFFFFFFFF +#define WARP_XOR(val, mask) __shfl_xor_sync(FINAL_MASK, (val), (mask), 32) +#endif __device__ __host__ __forceinline__ int ceil_div(int a, int b) { @@ -12,7 +20,7 @@ __device__ __forceinline__ float warpReduceMax(float val) { #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) - val = fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32)); + val = fmaxf(val, WARP_XOR(val, mask)); return val; } @@ -21,7 +29,7 @@ __device__ __forceinline__ float warpReduceMin(float val) { #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) - val = fminf(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32)); + val = fminf(val, WARP_XOR(val, mask)); return val; } @@ -345,7 +353,13 @@ fake_int4_quant_cuda( at::ScalarType::BFloat16, x.scalar_type(), "int4_quant_cuda", [&] { launch_int4_quant_kernel( +#if defined(__HIP_PLATFORM_AMD__) + // the templated const_data_ptr does not link under hipcc: clang mangles + // its enable_if template parameter differently from the gcc that built libtorch + static_cast(x.const_data_ptr()), +#else x.const_data_ptr(), +#endif out.data_ptr(), out_scale.data_ptr(), out_zero.data_ptr(), diff --git a/vime/backends/megatron_utils/kernels/int4_qat/setup.py b/vime/backends/megatron_utils/kernels/int4_qat/setup.py index 8715dd7b8..2db4683f9 100644 --- a/vime/backends/megatron_utils/kernels/int4_qat/setup.py +++ b/vime/backends/megatron_utils/kernels/int4_qat/setup.py @@ -3,6 +3,13 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension import torch +# A ROCm PyTorch build makes CUDAExtension hipify the sources and call hipcc, +# which rejects the nvcc-only flags below. The gfx target is passed through +# PYTORCH_ROCM_ARCH instead of -gencode. +IS_ROCM = torch.version.hip is not None +if IS_ROCM: + os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950") + # Get CUDA arch list arch_list = [] if torch.cuda.is_available(): @@ -32,18 +39,22 @@ "-O3", "-std=c++17", ], - "nvcc": [ - "-O3", - "-std=c++17", - "--expt-relaxed-constexpr", - "-Xcompiler", - "-fPIC", - ] - + [ - f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}' - for arch in arch_list - ] - + ["-gencode=arch=compute_90a,code=sm_90a"], + "nvcc": ( + ["-O3", "-std=c++17"] + if IS_ROCM + else [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + "-Xcompiler", + "-fPIC", + ] + + [ + f'-gencode=arch=compute_{arch.replace(".", "")},code=sm_{arch.replace(".", "")}' + for arch in arch_list + ] + + ["-gencode=arch=compute_90a,code=sm_90a"] + ), }, ) ], diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 3f6ab29c5..f4ad09898 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -13,6 +13,7 @@ from vime.utils.ppo_utils import ( calculate_log_probs_and_entropy, compute_approx_kl, + compute_cispo_loss, compute_gspo_kl, compute_opsm_mask, compute_policy_loss, @@ -30,6 +31,68 @@ slice_log_prob_with_cp, ) +ROLLOUT_TOP_P_TOKEN_KEYS = ( + "rollout_top_p_token_ids", + "rollout_top_p_token_offsets", +) + + +# Optional capture of per-sample policy log-probs computed during the training +# forward. Used only when dumping train debug data in configs that skip the +# separate log-prob recompute (can_reuse_log_probs_in_loss / use_rollout_logprobs): +# the values are identical to a separate compute_log_prob pass, so we snapshot +# them here at no extra forward. Keyed by GLOBAL rollout position so the writer +# can put them back in sample order regardless of pipeline/microbatch order. +_LOG_PROB_CAPTURE: "dict[int, torch.Tensor] | None" = None + + +def enable_log_prob_capture() -> None: + """Start capturing training-forward log-probs (call before ``train``).""" + global _LOG_PROB_CAPTURE + _LOG_PROB_CAPTURE = {} + + +def drain_captured_log_probs() -> "dict[int, torch.Tensor]": + """Return captured ``{rollout_position: cp-local log_probs}`` and stop capturing.""" + global _LOG_PROB_CAPTURE + captured = _LOG_PROB_CAPTURE or {} + _LOG_PROB_CAPTURE = None + return captured + + +def _maybe_capture_log_probs(batch: RolloutBatch, log_probs: list[torch.Tensor]) -> None: + """Snapshot per-sample CP-local ``log_probs`` keyed by global rollout position. + + No-op unless :func:`enable_log_prob_capture` is active and the micro-batch + carries ``partition`` (only added to the training keys when dumping). First + occurrence per position wins, so multi-step training keeps the initial + (old-policy) values. ``log_probs`` here is the per-sample list, in the same + order as ``batch['partition']`` (both indexed by this micro-batch's + ``micro_batch_indices``). + """ + if _LOG_PROB_CAPTURE is None: + return + positions = batch.get("partition") + if not positions: + return + for position, log_prob in zip(positions, log_probs, strict=True): + if position not in _LOG_PROB_CAPTURE: + _LOG_PROB_CAPTURE[position] = log_prob.detach().clone() + + +def get_rollout_top_p_logprob_kwargs(args: Namespace, batch: dict[str, Any]) -> dict[str, Any]: + if args.rollout_top_p == 1.0: + return {} + + top_p_token_ids = batch.get("rollout_top_p_token_ids") + top_p_token_offsets = batch.get("rollout_top_p_token_offsets") + if top_p_token_ids is None or top_p_token_offsets is None: + raise ValueError("rollout_top_p != 1.0 requires rollout_top_p_token_ids and rollout_top_p_token_offsets.") + return { + "top_p_token_ids": top_p_token_ids, + "top_p_token_offsets": top_p_token_offsets, + } + def get_responses( logits: torch.Tensor, @@ -38,7 +101,6 @@ def get_responses( unconcat_tokens: list[torch.Tensor], total_lengths: list[int], response_lengths: list[int], - max_seq_lens: list[int] | None = None, apply_temperature: bool = True, ) -> Iterator[tuple[torch.Tensor, torch.Tensor]]: """Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample. @@ -63,17 +125,10 @@ def get_responses( `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]` (1D int64), both aligned to response tokens for one sample. """ - qkv_format = args.qkv_format - assert logits.dtype == torch.float32, f"{logits.dtype}" assert len(logits.shape) == 3, f"{logits.shape}" - - if qkv_format == "thd": - assert logits.size(0) == 1, f"{logits.shape}" - logits = logits.squeeze(0) - else: - assert max_seq_lens is not None - logits = logits.view(-1, logits.size(-1)) + assert logits.size(0) == 1, f"{logits.shape}" + logits = logits.squeeze(0) if apply_temperature and args.rollout_temperature != 1.0: logits = logits.div(args.rollout_temperature) @@ -81,18 +136,10 @@ def get_responses( cp_size = mpu.get_context_parallel_world_size() end = 0 seq_start = 0 - for i, (tokens, total_length, response_length) in enumerate( - zip(unconcat_tokens, total_lengths, response_lengths, strict=False) - ): - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None - + for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False): if cp_size == 1: - if qkv_format == "bshd": - end = max_seq_len * i + total_length - start = end - response_length - else: - end += total_length - start = end - response_length + end += total_length + start = end - response_length logits_chunk = logits[start - 1 : end - 1] tokens_chunk = tokens[-response_length:] elif args.allgather_cp: @@ -121,7 +168,7 @@ def get_responses( else: # TODO: this is super ugly... do better abstraction. chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length, qkv_format, max_seq_len + total_length, response_length ) logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size] @@ -148,10 +195,8 @@ def _allgather_cp_redistribute( res: dict[str, list[torch.Tensor]], *, logits_local_len: int, - args: Namespace, total_lengths: list[int], response_lengths: list[int], - max_seq_lens: list[int] | None = None, ) -> None: """Redistribute response tensors from allgather-CP layout to zigzag ring-attn layout. @@ -165,10 +210,8 @@ def _allgather_cp_redistribute( Args: res: Dict mapping metric names to lists of per-sample tensors. logits_local_len: Local sequence length on this rank. - args: Configuration (needs ``qkv_format``). total_lengths: Total sequence lengths (prompt + response) per sample. response_lengths: Response segment lengths per sample. - max_seq_lens: Optional padded max sequence lengths per sample. """ cp_group = mpu.get_context_parallel_group() cp_rank = mpu.get_context_parallel_rank() @@ -202,7 +245,7 @@ def _allgather_cp_redistribute( response_length, dtype=ref_dtype, device=ref_device, - requires_grad=True, + requires_grad=ref_value.requires_grad, ) else: resp_start = s - logit_global_start @@ -213,19 +256,16 @@ def _allgather_cp_redistribute( full_resps.append(full_resp) seq_start += total_length - # Single differentiable all-reduce to gather full response from all CP ranks + # Single differentiable all-reduce to gather full response from all CP ranks. all_cat = torch.cat(full_resps, dim=0) all_cat = dist.nn.all_reduce(all_cat, group=cp_group) # Re-slice each sample into zigzag CP pattern new_values = [] - for idx, (full_resp, total_length, response_length) in enumerate( - zip(all_cat.split(response_lengths, dim=0), total_lengths, response_lengths, strict=False) + for full_resp, total_length, response_length in zip( + all_cat.split(response_lengths, dim=0), total_lengths, response_lengths, strict=False ): - max_seq_len = max_seq_lens[idx] if max_seq_lens is not None else None - new_values.append( - slice_log_prob_with_cp(full_resp, total_length, response_length, args.qkv_format, max_seq_len) - ) + new_values.append(slice_log_prob_with_cp(full_resp, total_length, response_length)) res[key] = new_values @@ -236,8 +276,6 @@ def _build_shifted_tokens( unconcat_tokens: list[torch.Tensor], total_lengths: list[int], response_lengths: list[int], - qkv_format: str, - max_seq_lens: list[int] | None, allgather_cp: bool, ) -> torch.Tensor: """Build shifted target tokens for the full packed/padded logits.""" @@ -247,12 +285,11 @@ def _build_shifted_tokens( if cp_size > 1 and not allgather_cp: full_tokens = torch.zeros(T, dtype=torch.long, device=device) end = 0 - for i, (tokens, total_length, response_length) in enumerate( - zip(unconcat_tokens, total_lengths, response_lengths, strict=False) + for tokens, total_length, response_length in zip( + unconcat_tokens, total_lengths, response_lengths, strict=False ): - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None chunk_size_cp, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length, qkv_format, max_seq_len + total_length, response_length ) for half, base in ((0, end), (1, end + chunk_size_cp)): lo = logits_offset[half][0] - chunks_offset[half][0] @@ -265,15 +302,10 @@ def _build_shifted_tokens( T_global = sum(total_lengths) if allgather_cp else T full_tokens = torch.zeros(T_global, dtype=torch.long, device=device) - if qkv_format == "thd" or allgather_cp: - offset = 0 - for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False): - full_tokens[offset : offset + total_length - 1] = tokens[1:total_length] - offset += total_length - else: # bshd, cp1 - for i, (tokens, total_length) in enumerate(zip(unconcat_tokens, total_lengths, strict=False)): - seq_start = max_seq_lens[i] * i - full_tokens[seq_start : seq_start + total_length - 1] = tokens[1:total_length] + offset = 0 + for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False): + full_tokens[offset : offset + total_length - 1] = tokens[1:total_length] + offset += total_length # allgather-CP: slice to local chunk if allgather_cp: @@ -291,13 +323,117 @@ def _build_shifted_tokens( return full_tokens +def _fill_topp_mask_rows( + keep: torch.Tensor, + ids: list[int], + offsets: list[int], + response_start: int, + local_start: int, + length: int, + vocab_start: int, + vocab_end: int, +) -> None: + end = min(response_start + length, max(len(offsets) - 1, 0)) + for response_idx in range(response_start, end): + local_ids = [ + token_id - vocab_start + for token_id in ids[offsets[response_idx] : offsets[response_idx + 1]] + if vocab_start <= token_id < vocab_end + ] + row = local_start + response_idx - response_start + keep[row].fill_(False) + if local_ids: + keep[row, torch.tensor(local_ids, device=keep.device, dtype=torch.long)] = True + + +def _build_topp_keep_mask( + T: int, + vocab_local: int, + device: torch.device, + top_p_token_ids: list[list[int]], + top_p_token_offsets: list[list[int]], + total_lengths: list[int], + response_lengths: list[int], + allgather_cp: bool, +) -> torch.Tensor: + """Build a ``[T, vocab_local]`` boolean keep-mask aligned to local logits. + + For response token ``r`` of a sample, the rollout top-p nucleus is + ``ids[offsets[r]:offsets[r + 1]]``. Rows without a recorded nucleus stay + all-True, so only response rows with replay data are masked. + """ + cp_size = mpu.get_context_parallel_world_size() + tp_rank = mpu.get_tensor_model_parallel_rank() + vocab_start = tp_rank * vocab_local + vocab_end = vocab_start + vocab_local + + # Normalize ragged payloads (may arrive as CPU int32 tensors) to python lists. + top_p_token_ids = [t.tolist() if torch.is_tensor(t) else list(t) for t in top_p_token_ids] + top_p_token_offsets = [t.tolist() if torch.is_tensor(t) else list(t) for t in top_p_token_offsets] + + keep = torch.ones((T, vocab_local), dtype=torch.bool, device=device) + + if cp_size > 1 and not allgather_cp: + local_base = 0 + for ids, offsets, total_length, response_length in zip( + top_p_token_ids, top_p_token_offsets, total_lengths, response_lengths, strict=False + ): + prompt_length = total_length - response_length + chunk_size_cp, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( + total_length, response_length + ) + for half, base in ((0, local_base), (1, local_base + chunk_size_cp)): + local_start = base + logits_offset[half][0] - chunks_offset[half][0] + length = logits_offset[half][1] - logits_offset[half][0] + response_start = tokens_offset[half][0] - prompt_length + _fill_topp_mask_rows(keep, ids, offsets, response_start, local_start, length, vocab_start, vocab_end) + local_base += 2 * chunk_size_cp + return keep + + if allgather_cp: + cp_rank = mpu.get_context_parallel_rank() + chunk_start = cp_rank * T + chunk_end = chunk_start + T + seq_start = 0 + for ids, offsets, total_length, response_length in zip( + top_p_token_ids, top_p_token_offsets, total_lengths, response_lengths, strict=False + ): + prompt_length = total_length - response_length + logit_global_start = seq_start + prompt_length - 1 + logit_global_end = seq_start + total_length - 1 + s = max(logit_global_start, chunk_start) + e = min(logit_global_end, chunk_end) + if e > s: + _fill_topp_mask_rows( + keep, + ids, + offsets, + s - logit_global_start, + s - chunk_start, + e - s, + vocab_start, + vocab_end, + ) + seq_start += total_length + return keep + + offset = 0 + for ids, offsets, total_length, response_length in zip( + top_p_token_ids, top_p_token_offsets, total_lengths, response_lengths, strict=False + ): + end = offset + total_length + start = end - response_length + _fill_topp_mask_rows(keep, ids, offsets, 0, start - 1, response_length, vocab_start, vocab_end) + offset += total_length + + return keep + + def _extract_per_sample( log_prob_full: torch.Tensor, entropy_full: torch.Tensor | None, total_lengths: list[int], response_lengths: list[int], - qkv_format: str, - max_seq_lens: list[int] | None, allgather_cp: bool, ) -> tuple[list[torch.Tensor], list[torch.Tensor | None]]: """Slice per-sample response log-probs/entropy from full-length 1-D tensors.""" @@ -308,10 +444,9 @@ def _extract_per_sample( if cp_size > 1 and not allgather_cp: # zigzag CP pos = 0 - for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)): - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + for total_length, response_length in zip(total_lengths, response_lengths, strict=False): chunk_size_cp, chunks_offset, logits_offset, _tokens_offset = get_logits_and_tokens_offset_with_cp( - total_length, response_length, qkv_format, max_seq_len + total_length, response_length ) lo0 = logits_offset[0][0] - chunks_offset[0][0] hi0 = logits_offset[0][1] - chunks_offset[0][0] @@ -352,9 +487,9 @@ def _extract_per_sample( s = max(logit_global_start, chunk_start) e = min(logit_global_end, chunk_end) if e <= s: - log_probs_list.append(torch.zeros((0,), dtype=log_prob_full.dtype, device=log_prob_full.device)) + log_probs_list.append(log_prob_full[:0]) if entropy_full is not None: - entropy_list.append(torch.zeros((0,), dtype=entropy_full.dtype, device=entropy_full.device)) + entropy_list.append(entropy_full[:0]) else: log_probs_list.append(log_prob_full[s - chunk_start : e - chunk_start]) if entropy_full is not None: @@ -363,22 +498,14 @@ def _extract_per_sample( else: # cp1 - if qkv_format == "thd": - offset = 0 - for total_length, response_length in zip(total_lengths, response_lengths, strict=False): - end = offset + total_length - start = end - response_length - log_probs_list.append(log_prob_full[start - 1 : end - 1]) - if entropy_full is not None: - entropy_list.append(entropy_full[start - 1 : end - 1]) - offset += total_length - else: # bshd - for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)): - end = max_seq_lens[i] * i + total_length - start = end - response_length - log_probs_list.append(log_prob_full[start - 1 : end - 1]) - if entropy_full is not None: - entropy_list.append(entropy_full[start - 1 : end - 1]) + offset = 0 + for total_length, response_length in zip(total_lengths, response_lengths, strict=False): + end = offset + total_length + start = end - response_length + log_probs_list.append(log_prob_full[start - 1 : end - 1]) + if entropy_full is not None: + entropy_list.append(entropy_full[start - 1 : end - 1]) + offset += total_length return log_probs_list, entropy_list @@ -392,7 +519,8 @@ def get_log_probs_and_entropy( response_lengths: list[int], with_entropy: bool = False, non_loss_data: bool = True, - max_seq_lens: list[int] | None = None, + top_p_token_ids: list[list[int]] | None = None, + top_p_token_offsets: list[list[int]] | None = None, ) -> dict[str, list[torch.Tensor]]: """Compute per-token log-probabilities (and optionally entropy) on responses. @@ -400,21 +528,14 @@ def get_log_probs_and_entropy( per-sample slicing) so backward traverses ``[T, V]`` only once, then extracts per-sample response portions. - When ``entropy_coef == 0``, entropy is computed under ``torch.no_grad()`` - to avoid retaining the computation graph and to skip cloning. + If rollout top-p replay is provided, the keep-mask is applied only to + log-probabilities; entropy is always computed from the unmasked logits. """ assert non_loss_data - qkv_format = args.qkv_format - assert logits.dtype == torch.float32, f"{logits.dtype}" assert len(logits.shape) == 3, f"{logits.shape}" - - if qkv_format == "thd": - assert logits.size(0) == 1, f"{logits.shape}" - logits = logits.squeeze(0) - else: - assert max_seq_lens is not None - logits = logits.view(-1, logits.size(-1)) + assert logits.size(0) == 1, f"{logits.shape}" + logits = logits.squeeze(0) # Apply rollout temperature scaling to logits to match rollout-time log-probs. rollout_temperature = getattr(args, "rollout_temperature", 1.0) @@ -425,11 +546,26 @@ def get_log_probs_and_entropy( device = logits.device tp_group = mpu.get_tensor_model_parallel_group() chunk_size = args.log_probs_chunk_size + # Keep entropy metrics, but skip saving entropy-backward activations when + # the entropy term cannot affect the loss. + with_entropy_grad = with_entropy and getattr(args, "entropy_coef", 0.0) != 0 # --- build full shifted-token target tensor --- - full_tokens = _build_shifted_tokens( - T, device, unconcat_tokens, total_lengths, response_lengths, qkv_format, max_seq_lens, args.allgather_cp - ) + full_tokens = _build_shifted_tokens(T, device, unconcat_tokens, total_lengths, response_lengths, args.allgather_cp) + + # --- build top-p nucleus keep-mask (logprob only; entropy stays unmasked) --- + top_p_keep_mask = None + if top_p_token_ids is not None and top_p_token_offsets is not None: + top_p_keep_mask = _build_topp_keep_mask( + T, + logits.size(-1), + device, + top_p_token_ids, + top_p_token_offsets, + total_lengths, + response_lengths, + args.allgather_cp, + ) # --- compute on full [T,V] logits at once via calculate_log_probs_and_entropy --- log_prob_full, entropy_full = calculate_log_probs_and_entropy( @@ -437,7 +573,9 @@ def get_log_probs_and_entropy( full_tokens, tp_group, with_entropy=with_entropy, + with_entropy_grad=with_entropy_grad, chunk_size=chunk_size, + log_prob_keep_mask=top_p_keep_mask, ) log_prob_full = log_prob_full.squeeze(-1) # [T, 1] -> [T] @@ -447,8 +585,6 @@ def get_log_probs_and_entropy( entropy_full, total_lengths, response_lengths, - qkv_format, - max_seq_lens, args.allgather_cp, ) @@ -461,10 +597,8 @@ def get_log_probs_and_entropy( _allgather_cp_redistribute( res, logits_local_len=T, - args=args, total_lengths=total_lengths, response_lengths=response_lengths, - max_seq_lens=max_seq_lens, ) return torch.empty((0,), device=device), res @@ -479,7 +613,6 @@ def get_values( response_lengths: list[int], with_entropy: bool = False, non_loss_data: bool = True, - max_seq_lens: list[int] | None = None, ) -> dict[str, list[torch.Tensor]]: """Extract per-token value predictions over response tokens. @@ -507,7 +640,6 @@ def get_values( unconcat_tokens=unconcat_tokens, total_lengths=total_lengths, response_lengths=response_lengths, - max_seq_lens=max_seq_lens, apply_temperature=False, ): assert logits_chunk.size(-1) == 1, f"{logits_chunk.shape}" @@ -521,10 +653,8 @@ def get_values( _allgather_cp_redistribute( res, logits_local_len=logits.size(1), - args=args, total_lengths=total_lengths, response_lengths=response_lengths, - max_seq_lens=max_seq_lens, ) return torch.empty((0,), device=logits.device), res @@ -576,10 +706,10 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) This function extracts rewards, log-probs, values, and masks from `rollout_data`, computes KL divergences, then applies the chosen advantage - estimator. Supported methods: "grpo", "gspo", "ppo", "reinforce_plus_plus", - and "reinforce_plus_plus_baseline". When `args.normalize_advantages` is - True, advantages are whitened across the data-parallel group using masked - statistics. + estimator. Supported methods: "grpo", "gspo", "cispo", "ppo", + "reinforce_plus_plus", and "reinforce_plus_plus_baseline". When + `args.normalize_advantages` is True, advantages are whitened across the + data-parallel-with-context-parallel group using masked statistics. Early returns if both `log_probs` and `values` are None (intermediate pipeline stages). @@ -606,8 +736,6 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) response_lengths: list[int] = rollout_data.get("response_lengths") loss_masks: list[torch.Tensor] = rollout_data.get("loss_masks") total_lengths: list[int] = rollout_data.get("total_lengths") - max_seq_lens: list[int] | None = rollout_data.get("max_seq_lens", None) - # return when not the last pp stage. if not mpu.is_pipeline_last_stage(): return @@ -632,7 +760,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) custom_adv_fn(args, rollout_data) advantages, returns = rollout_data["advantages"], rollout_data["returns"] - elif args.advantage_estimator in ["grpo", "gspo"]: + elif args.advantage_estimator in ["grpo", "gspo", "cispo"]: rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) returns = get_grpo_returns(rewards, kl) # TODO: is the copy necessary? @@ -643,11 +771,11 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) rewards = [] kl_coef = -args.kl_coef cp_rank = mpu.get_context_parallel_rank() - for reward, k in zip(old_rewards, kl, strict=False): - k *= kl_coef + for reward, per_token_kl in zip(old_rewards, kl, strict=False): + token_level_rewards = per_token_kl * kl_coef if cp_rank == 0: - k[-1] += reward - rewards.append(k) + token_level_rewards[-1] += reward + rewards.append(token_level_rewards) advantages, returns = get_advantages_and_returns_batch( total_lengths, response_lengths, values, rewards, args.gamma, args.lambd ) @@ -670,7 +798,6 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) advantages = get_reinforce_plus_plus_baseline_advantages( rewards=rewards, kl=kl, - loss_masks=loss_masks, kl_coef=args.kl_coef, ) returns = advantages @@ -699,11 +826,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) total_len = total_lengths[i] response_len = response_lengths[i] prompt_len = total_len - response_len - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None - - _, _, _, token_offsets = get_logits_and_tokens_offset_with_cp( - total_len, response_len, args.qkv_format, max_seq_len - ) + _, _, _, token_offsets = get_logits_and_tokens_offset_with_cp(total_len, response_len) # Convert global offsets to response-space offsets s0, e0 = token_offsets[0] @@ -728,20 +851,30 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) all_masks = torch.cat(mask_chunks) - if all_masks.numel() > 0: - assert ( - all_advs.size() == all_masks.size() - ), f"Shape mismatch before whitening: advantages {all_advs.size()}, masks {all_masks.size()}" - dp_group = mpu.get_data_parallel_group() - - whitened_advs_flat = distributed_masked_whiten( - all_advs, - all_masks, - process_group=dp_group, - shift_mean=True, - ) - chunk_lengths = [chunk.size(0) for chunk in advantages] - advantages = list(torch.split(whitened_advs_flat, chunk_lengths)) + assert ( + all_advs.size() == all_masks.size() + ), f"Shape mismatch before whitening: advantages {all_advs.size()}, masks {all_masks.size()}" + # `all_advs` / `all_masks` only cover the tokens this CP rank owns, so the + # statistics must be reduced over the DP group *with* context parallel. + # The CP-excluding group makes every CP rank whiten its own zigzag slice + # with its own mean/var, i.e. the two halves of one sequence get + # different affine transforms. + # + # This has to stay unconditional: a CP rank can legitimately own zero + # response tokens (prompt-heavy sequences put both of its chunks inside + # the prompt), and skipping the collective on just that rank would + # desync the all_reduce. `distributed_masked_whiten` handles an empty + # local tensor — it contributes 0 to the reduced sums. + dp_cp_group = mpu.get_data_parallel_group(with_context_parallel=True) + + whitened_advs_flat = distributed_masked_whiten( + all_advs, + all_masks, + process_group=dp_cp_group, + shift_mean=True, + ) + chunk_lengths = [chunk.size(0) for chunk in advantages] + advantages = list(torch.split(whitened_advs_flat, chunk_lengths)) rollout_data["advantages"] = advantages rollout_data["returns"] = returns @@ -832,7 +965,6 @@ def policy_loss_function( response_lengths = batch["response_lengths"] total_lengths = batch["total_lengths"] - max_seq_lens = batch.get("max_seq_lens", None) _, log_probs_and_entropy = get_log_probs_and_entropy( logits, @@ -841,10 +973,14 @@ def policy_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=True, - max_seq_lens=max_seq_lens, + **get_rollout_top_p_logprob_kwargs(args, batch), ) log_probs = log_probs_and_entropy["log_probs"] + # Snapshot the per-sample policy log-probs for the train debug dump (no-op + # unless capture is enabled). Must run before the torch.cat below rebinds + # `log_probs` to a single concatenated tensor. + _maybe_capture_log_probs(batch, log_probs) if not args.use_rollout_logprobs and not old_log_probs: old_log_probs = [log_prob.detach() for log_prob in log_probs] train_log_probs_for_tis = batch.get("log_probs") @@ -895,7 +1031,16 @@ def policy_loss_function( log_probs = torch.cat(log_probs, dim=0) ppo_kl = old_log_probs - log_probs - pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) + if args.advantage_estimator == "cispo": + pg_loss, pg_clipfrac = compute_cispo_loss(ppo_kl, log_probs, advantages, args.eps_clip, args.eps_clip_high) + else: + pg_loss, pg_clipfrac = compute_policy_loss( + ppo_kl, + advantages, + args.eps_clip, + args.eps_clip_high, + eps_clip_c=args.eps_clip_c, + ) if args.use_opsm: pg_loss = pg_loss * opsm_mask @@ -943,8 +1088,6 @@ def policy_loss_function( modified_response_masks, batch["rollout_mask_sums"], args.calculate_per_token_loss, - args.qkv_format, - max_seq_lens, ) # Determine pg_loss reducer: use custom if specified, otherwise default @@ -992,7 +1135,8 @@ def policy_loss_function( train_rollout_logprob_abs_diff = None if "rollout_log_probs" in batch and batch["rollout_log_probs"]: rollout_log_probs = torch.cat(batch["rollout_log_probs"], dim=0) - train_rollout_logprob_abs_diff = sum_of_sample_mean((old_log_probs - rollout_log_probs).abs()) + log_probs_to_compare = log_probs if args.use_rollout_logprobs else old_log_probs + train_rollout_logprob_abs_diff = sum_of_sample_mean((log_probs_to_compare - rollout_log_probs).abs()) reported_loss = { "loss": loss.clone().detach(), @@ -1021,7 +1165,7 @@ def policy_loss_function( reported_loss["opsm_clipfrac"] = opsm_clipfrac # Add OPD metrics if available - if "opd_reverse_kl" in batch: + if batch.get("opd_reverse_kl"): opd_reverse_kl = torch.cat(batch["opd_reverse_kl"], dim=0) reported_loss["opd_reverse_kl"] = sum_of_sample_mean(opd_reverse_kl).clone().detach() @@ -1059,7 +1203,6 @@ def value_loss_function( unconcat_tokens=batch["unconcat_tokens"], total_lengths=batch["total_lengths"], response_lengths=batch["response_lengths"], - max_seq_lens=batch.get("max_seq_lens", None), ) values = torch.cat([value.flatten() for value in values["values"]], dim=0) @@ -1118,7 +1261,6 @@ def sft_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=False, - max_seq_lens=batch.get("max_seq_lens", None), ) log_probs = log_probs_and_entropy["log_probs"] @@ -1179,8 +1321,6 @@ def loss_function( batch["loss_masks"], batch["rollout_mask_sums"], args.calculate_per_token_loss, - args.qkv_format, - batch.get("max_seq_lens", None), ) match args.loss_type: diff --git a/vime/backends/megatron_utils/megatron_patch/__init__.py b/vime/backends/megatron_utils/megatron_patch/__init__.py deleted file mode 100644 index 8693ff85a..000000000 --- a/vime/backends/megatron_utils/megatron_patch/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from . import megatron_chunked_grad_coalesce_patch # noqa: F401 diff --git a/vime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py b/vime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py deleted file mode 100644 index 800da9b0c..000000000 --- a/vime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py +++ /dev/null @@ -1,146 +0,0 @@ -# Patch _allreduce_non_tensor_model_parallel_grads (and its legacy alias -# _allreduce_layernorm_grads) in megatron.core.distributed.finalize_model_grads -# to coalesce/all_reduce TP-side grads in size-bounded chunks instead of one -# large _flatten_dense_tensors(grads). Lowers the peak contiguous-memory -# allocation during TP-side grad sync, avoiding OOM under allocator -# fragmentation when the combined grad buffer would otherwise be very large. -# SUM/AVG are element-wise, so chunking is mathematically equivalent. -# Chunk size: VIME_GRAD_COALESCE_CHUNK_BYTES, default 1 GiB. -# -# Cross-compatible across the Megatron-LM versions vime is run against: -# the core_v0.13.0 line (DDP config exposes `use_custom_fsdp`, `_get_main_grad_attr` -# takes `(param, use_custom_fsdp)`, target function takes `(model, config)`) and -# the post-core_v0.15.0rc7 dev line (`use_megatron_fsdp`, single-arg -# `_get_main_grad_attr`, `(model, config, tp_group)`). API differences are -# resolved at runtime — no version-conditional imports. - -import inspect -import logging -import os -import sys -import warnings - -logger = logging.getLogger(__name__) - -try: - import torch - from megatron.core import parallel_state - from megatron.core.distributed.finalize_model_grads import ( - _flatten_dense_tensors, - _get_main_grad_attr, - _reshard_if_dtensor, - _unflatten_dense_tensors, - _unshard_if_dtensor, - get_attr_wrapped_model, - ) - - # post-core_v0.15.0rc7 dev takes (param); core_v0.13.0 line takes - # (param, use_custom_fsdp=False). - _gma_takes_fsdp_arg = len(inspect.signature(_get_main_grad_attr).parameters) >= 2 - - def _grad_attr(param, fsdp_on): - if _gma_takes_fsdp_arg: - return _get_main_grad_attr(param, fsdp_on) - return _get_main_grad_attr(param) - - def _fsdp_flag(ddp_config): - return bool(getattr(ddp_config, "use_megatron_fsdp", False) or getattr(ddp_config, "use_custom_fsdp", False)) - - _chunk_bytes = int(os.environ.get("VIME_GRAD_COALESCE_CHUNK_BYTES") or (1 << 30)) - - def _split_into_chunks(params, grads, target_bytes): - """Greedy split keeping params/grads aligned. A single grad larger - than target_bytes is placed alone in its own chunk.""" - chunks, cur_p, cur_g, cur_b = [], [], [], 0 - for p, g in zip(params, grads, strict=False): - gb = g.numel() * g.element_size() - if cur_g and cur_b + gb > target_bytes: - chunks.append((cur_p, cur_g)) - cur_p, cur_g, cur_b = [], [], 0 - cur_p.append(p) - cur_g.append(g) - cur_b += gb - if cur_g: - chunks.append((cur_p, cur_g)) - return chunks - - def _allreduce_non_tensor_model_parallel_grads(model, config, tp_group=None): - # post-core_v0.15.0rc7 dev passes tp_group; core_v0.13.0 line omits it. - # Default-fill from parallel_state so the same body works for both call sites. - if tp_group is None: - tp_group = parallel_state.get_tensor_model_parallel_group() - if tp_group.size() <= 1: - return - - params_sum, grads_sum = [], [] - params_avg, grads_avg = [], [] - ddp_config = None - for model_chunk in model: - ddp_config = model_chunk.ddp_config - fsdp_on = _fsdp_flag(ddp_config) - for name, param in get_attr_wrapped_model(model_chunk, "named_parameters")(): - if not param.requires_grad: - continue - if getattr(param, "average_gradients_across_tp_domain", False): - target_params, target_grads = params_avg, grads_avg - elif (config.sequence_parallel and getattr(param, "sequence_parallel", False)) or ( - config.qk_layernorm and ("q_layernorm" in name or "k_layernorm" in name) - ): - target_params, target_grads = params_sum, grads_sum - else: - continue - - grad_attr = _grad_attr(param, fsdp_on) - grad = getattr(param, grad_attr) - if grad is None: - continue - target_params.append(param) - if fsdp_on and hasattr(grad, "_local_tensor"): - target_grads.append(grad._local_tensor.data) - else: - target_grads.append(_unshard_if_dtensor(grad).data) - - for params, grads, op in ( - (params_sum, grads_sum, torch.distributed.ReduceOp.SUM), - (params_avg, grads_avg, torch.distributed.ReduceOp.AVG), - ): - if not grads: - continue - fsdp_on = _fsdp_flag(ddp_config) - for p_chunk, g_chunk in _split_into_chunks(params, grads, _chunk_bytes): - coalesced = _flatten_dense_tensors(g_chunk) - torch.distributed.all_reduce(coalesced, op=op, group=tp_group) - for param, buf, synced in zip( - p_chunk, g_chunk, _unflatten_dense_tensors(coalesced, g_chunk), strict=False - ): - buf.copy_(synced) - grad_attr = _grad_attr(param, fsdp_on) - orig_grad = getattr(param, grad_attr) - if fsdp_on and hasattr(orig_grad, "_local_tensor"): - # buf already aliases orig_grad._local_tensor.data; - # restore original DTensor wrapper (post-rc7 dev semantics). - setattr(param, grad_attr, orig_grad) - else: - setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad)) - del coalesced - - # The parent package re-exports a same-named function, shadowing the - # submodule attribute. Pull the real module out of sys.modules to setattr. - _fmg = sys.modules["megatron.core.distributed.finalize_model_grads"] - _fmg._allreduce_non_tensor_model_parallel_grads = _allreduce_non_tensor_model_parallel_grads - _fmg._allreduce_layernorm_grads = _allreduce_non_tensor_model_parallel_grads - - logger.info( - "vime grad coalesce patch applied to " - "megatron.core.distributed.finalize_model_grads." - "_allreduce_non_tensor_model_parallel_grads (chunk=%d MiB)", - _chunk_bytes // (1 << 20), - ) - -except ImportError as exc: - warnings.warn( - f"vime grad coalesce patch not applied — Megatron import failed ({exc!r}). " - "If this is a Megatron upgrade, the symbol layout may have changed; " - "without this patch, large-model TP grad sync may OOM.", - stacklevel=2, - ) diff --git a/vime/backends/megatron_utils/megatron_to_hf/__init__.py b/vime/backends/megatron_utils/megatron_to_hf/__init__.py index f14217862..a06446e68 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/vime/backends/megatron_utils/megatron_to_hf/__init__.py @@ -1,7 +1,6 @@ from .deepseekv3 import convert_deepseekv3_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf -from .gpt_oss import convert_gpt_oss_to_hf from .llama import convert_llama_to_hf from .mimo import convert_mimo_to_hf from .minimax_m2 import convert_minimax_m2_to_hf @@ -9,46 +8,48 @@ from .qwen2 import convert_qwen2_to_hf from .qwen3_5 import convert_qwen3_5_to_hf from .qwen3_next import convert_qwen3_next_to_hf +from .qwen3_omni import convert_qwen3_omni_to_hf from .qwen3_vl import convert_qwen3vl_to_hf from .qwen3moe import convert_qwen3moe_to_hf -# TODO unify w/ `convert_to_hf` -def postprocess_hf_param(args, megatron_param_name, hf_param_name, param): - param = remove_padding(megatron_param_name, param, args.vocab_size) - # TODO support quant - return param - - # TODO optimize code details -def convert_to_hf(args, model_name, name, param, quantization_config=None): - param = remove_padding(name, param, args.vocab_size) +def convert_to_hf(args, model_name, name, param, quantization_config=None, transform_ue8m0=False): + hf_name = name + while hf_name.startswith("module."): + hf_name = hf_name.removeprefix("module.") + if hf_name.startswith("model.visual."): + return [(hf_name, param)] + param = remove_padding(name, param, args.vocab_size) converted_named_tensors = _convert_to_hf_core(args, model_name, name, param) - return quantize_params(args, name, converted_named_tensors, quantization_config) + return quantize_params(args, name, converted_named_tensors, quantization_config, transform_ue8m0) # TODO optimize code details def _convert_to_hf_core(args, model_name, name, param): - if "minimaxm2" in model_name or "minimax_m2" in model_name: + model_name = model_name.lower().replace("_", "").replace("-", "") + if "minimaxm2" in model_name: converted_named_tensors = convert_minimax_m2_to_hf(args, name, param) - elif "glm4moelite" in model_name or "deepseekv3" in model_name: + elif any(family in model_name for family in ("glm4moelite", "deepseekv3", "deepseekv32", "glmmoedsa", "kimi")): converted_named_tensors = convert_deepseekv3_to_hf(args, name, param) elif "glm4moe" in model_name: converted_named_tensors = convert_glm4moe_to_hf(args, name, param) elif "glm4" in model_name: converted_named_tensors = convert_glm4_to_hf(args, name, param) - elif "gpt_oss" in model_name or "gpt-oss" in model_name or "gptoss" in model_name: - converted_named_tensors = convert_gpt_oss_to_hf(args, name, param) elif "qwen3moe" in model_name: converted_named_tensors = convert_qwen3moe_to_hf(args, name, param) + elif "qwen3omni" in model_name: + converted_named_tensors = convert_qwen3_omni_to_hf(args, name, param) elif "qwen3next" in model_name: converted_named_tensors = convert_qwen3_next_to_hf(args, name, param) - elif "qwen3_5" in model_name: + elif "qwen35" in model_name: converted_named_tensors = convert_qwen3_5_to_hf(args, name, param) elif "qwen3vl" in model_name: converted_named_tensors = convert_qwen3vl_to_hf(args, name, param) + elif "qwen2moe" in model_name or "qwen3moe" in model_name: + converted_named_tensors = convert_qwen3moe_to_hf(args, name, param) elif "qwen2" in model_name or "qwen3" in model_name: converted_named_tensors = convert_qwen2_to_hf(args, name, param) elif "llama" in model_name: diff --git a/vime/backends/megatron_utils/megatron_to_hf/gpt_oss.py b/vime/backends/megatron_utils/megatron_to_hf/gpt_oss.py deleted file mode 100644 index 0db77f912..000000000 --- a/vime/backends/megatron_utils/megatron_to_hf/gpt_oss.py +++ /dev/null @@ -1,105 +0,0 @@ -import re - -import torch - - -def convert_gpt_oss_to_hf(args, name, param): - """Convert Megatron GPT-OSS parameter names to HF format for weight update to vLLM.""" - - if name == "module.module.embedding.word_embeddings.weight": - return [("model.embed_tokens.weight", param)] - if name == "module.module.output_layer.weight": - return [("lm_head.weight", param)] - if name == "module.module.decoder.final_layernorm.weight": - return [("model.norm.weight", param)] - - head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads - value_num_per_group = args.num_attention_heads // args.num_query_groups - - decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" - match = re.match(decoder_layers_pattern, name) - if match: - layer_idx, rest = match.groups() - - # Expert weights - expert_pattern = r"mlp\.experts\.(.+)\.weight(\d+)" - match = re.match(expert_pattern, rest) - if match: - rest, expert_idx = match.groups() - if rest == "linear_fc1": - gate_weight, up_weight = param.chunk(2, dim=0) - return [ - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight), - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight), - ] - elif rest == "linear_fc2": - return [ - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param), - ] - else: - raise ValueError(f"Unknown expert parameter name: {name}") - - # Expert biases - expert_bias_pattern = r"mlp\.experts\.(.+)\.bias(\d+)" - match = re.match(expert_bias_pattern, rest) - if match: - rest, expert_idx = match.groups() - if rest == "linear_fc1": - gate_bias, up_bias = param.chunk(2, dim=0) - return [ - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.bias", gate_bias), - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.bias", up_bias), - ] - elif rest == "linear_fc2": - return [ - (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.bias", param), - ] - else: - raise ValueError(f"Unknown expert bias parameter name: {name}") - - # Attention - if rest == "self_attention.linear_proj.weight": - return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)] - elif rest == "self_attention.linear_proj.bias": - return [(f"model.layers.{layer_idx}.self_attn.o_proj.bias", param)] - elif rest == "self_attention.linear_qkv.weight": - param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size) - q_param, k_param, v_param = torch.split(param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1) - q_param = q_param.reshape(-1, args.hidden_size) - k_param = k_param.reshape(-1, args.hidden_size) - v_param = v_param.reshape(-1, args.hidden_size) - return [ - (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param), - (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param), - (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param), - ] - elif rest == "self_attention.linear_qkv.bias": - param = param.view(args.num_query_groups, -1) - q_bias, k_bias, v_bias = torch.split( - param, - split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim], - dim=1, - ) - q_bias = q_bias.contiguous().flatten() - k_bias = k_bias.contiguous().flatten() - v_bias = v_bias.contiguous().flatten() - return [ - (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias), - (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias), - (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias), - ] - # Learnable softmax offset (sinks) - elif rest == "self_attention.core_attention.softmax_offset": - return [(f"model.layers.{layer_idx}.self_attn.sinks", param)] - # Layer norms - elif rest == "self_attention.linear_qkv.layer_norm_weight": - return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] - elif rest == "pre_mlp_layernorm.weight": - return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] - # Router - elif rest == "mlp.router.weight": - return [(f"model.layers.{layer_idx}.mlp.router.weight", param)] - elif rest == "mlp.router.bias": - return [(f"model.layers.{layer_idx}.mlp.router.bias", param)] - - raise ValueError(f"Unknown parameter name: {name}") diff --git a/vime/backends/megatron_utils/megatron_to_hf/mimo.py b/vime/backends/megatron_utils/megatron_to_hf/mimo.py index 3d9c6c491..5efafb29b 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/mimo.py +++ b/vime/backends/megatron_utils/megatron_to_hf/mimo.py @@ -26,7 +26,7 @@ def convert_mimo_mtp_param(args, name, param): - Self attention (reuses Qwen2 attention structure) - MLP (reuses Qwen2 MLP structure) - Based on MimoBridge._convert_mtp_param logic (reverse mapping) + This is the inverse of MiMo's HuggingFace-to-Megatron MTP mapping. """ mtp_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)" match = re.match(mtp_pattern, name) @@ -37,7 +37,6 @@ def convert_mimo_mtp_param(args, name, param): layer_idx, component = match.groups() # Direct mappings for MTP-specific components (Megatron -> HF) - # Based on MimoBridge direct_name_mapping (reversed) direct_mappings = { "enorm.weight": f"model.mtp_layers.{layer_idx}.token_layernorm.weight", "hnorm.weight": f"model.mtp_layers.{layer_idx}.hidden_layernorm.weight", diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/vime/backends/megatron_utils/megatron_to_hf/processors/__init__.py index 64d106189..6efe6bdad 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/processors/__init__.py +++ b/vime/backends/megatron_utils/megatron_to_hf/processors/__init__.py @@ -1,18 +1,22 @@ from .padding_remover import remove_padding -from .quantizer_compressed_tensors import quantize_params_compressed_tensors -from .quantizer_fp8 import quantize_params_fp8 -__all__ = ["remove_padding", "quantize_param", "quantize_params_fp8", "quantize_params_compressed_tensors"] +__all__ = ["quantize_params", "remove_padding"] -def quantize_params(args, megatron_name, converted_named_params, quantization_config): +def quantize_params(args, megatron_name, converted_named_params, quantization_config, transform_ue8m0=True): if quantization_config is None: return converted_named_params - elif quantization_config["quant_method"] == "fp8": - return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config["quant_method"] == "compressed-tensors": + + if quantization_config["quant_method"] == "fp8": + from .quantizer_fp8 import quantize_params_fp8 + + return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config, transform_ue8m0) + + if quantization_config["quant_method"] == "compressed-tensors": + from .quantizer_compressed_tensors import quantize_params_compressed_tensors + # only int4 at the moment. return quantize_params_compressed_tensors(converted_named_params, quantization_config) - else: - # Unknown quant method (e.g. mxfp4) — pass through BF16 params as-is - return converted_named_params + + # Unknown quant method (e.g. mxfp4) — pass through BF16 params as-is + return converted_named_params diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py b/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py index 741584fa4..4e548adc9 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py +++ b/vime/backends/megatron_utils/megatron_to_hf/processors/padding_remover.py @@ -7,6 +7,7 @@ def remove_padding(name: str, param: torch.Tensor, vocab_size: int) -> torch.Ten """ Remove vocab padding: param[:vocab_size] for embedding/output layers, else unchanged. """ - if strip_param_name_prefix(name) in {"embedding.word_embeddings.weight", "output_layer.weight"}: + name = strip_param_name_prefix(name).removeprefix("language_model.") + if name in {"embedding.word_embeddings.weight", "output_layer.weight"}: return param[:vocab_size] return param diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py index ca69df8e6..b7f05a658 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py +++ b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py @@ -5,6 +5,8 @@ import torch import torch.nn as nn +from vime.utils import accelerator + try: import fake_int4_quant_cuda except ImportError: @@ -90,7 +92,7 @@ def from_linear(cls, linear, w_bit, group_size, init_only=False, scales=None, ze awq_linear.bias = linear.bias.clone().half() pack_num = 32 // awq_linear.w_bit - device = torch.device(f"cuda:{torch.cuda.current_device()}") + device = accelerator.current_device() repeat_scales = scales.to(device).t().repeat_interleave(group_size, 1) if isinstance(zeros, torch.Tensor): @@ -283,7 +285,7 @@ def quantize_params_compressed_tensors(converted_named_params, quantization_conf qw, s, zp = pack_layer(param, group_size, is_symmetric) qweight_name = name.replace(".weight", ".weight_packed") scale_name = name.replace(".weight", ".weight_scale") - weight_shape = torch.tensor(param.shape, dtype=torch.int32, device="cuda") + weight_shape = torch.tensor(param.shape, dtype=torch.int32, device=accelerator.device()) weight_shape_name = name.replace(".weight", ".weight_shape") if zp is not None: zp_name = name.replace(".weight", ".weight_zero_point") diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index 6d49f754d..25c9d720d 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -4,15 +4,16 @@ from vime.backends.megatron_utils.kernels.fp8_kernel import blockwise_cast_to_fp8_triton -from ...fp8_helpers import quant_weight_ue8m0, should_deepgemm_weight_requant_ue8m0 +from ...vllm import quant_weight_ue8m0, should_deepgemm_weight_requant_ue8m0, transform_scale_ue8m0 -def quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config): +def quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config, transform_ue8m0=True): assert quantization_config["quant_method"] == "fp8" fmt = quantization_config.get("fmt", "e4m3") assert fmt == "e4m3", f"Unsupported FP8 format: {fmt}" assert quantization_config["activation_scheme"] == "dynamic" weight_block_size = quantization_config.get("weight_block_size", None) + force_ue8m0_scale = getattr(args, "force_fp8_ue8m0_scale", False) decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" match = re.match(decoder_layers_pattern, megatron_name) @@ -43,7 +44,15 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio # TODO: find a clearer way. if converted_name.endswith("_scale"): continue - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size)) + quantize_named_params.extend( + _quantize_param( + converted_name, + param, + weight_block_size, + transform_ue8m0, + force_ue8m0_scale=force_ue8m0_scale, + ) + ) return quantize_named_params @@ -58,7 +67,15 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio ]: quantize_named_params = [] for converted_name, param in converted_named_params: - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size)) + quantize_named_params.extend( + _quantize_param( + converted_name, + param, + weight_block_size, + transform_ue8m0, + force_ue8m0_scale=force_ue8m0_scale, + ) + ) return quantize_named_params @@ -83,7 +100,15 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio ]: quantize_named_params = [] for converted_name, param in converted_named_params: - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size)) + quantize_named_params.extend( + _quantize_param( + converted_name, + param, + weight_block_size, + transform_ue8m0, + force_ue8m0_scale=force_ue8m0_scale, + ) + ) return quantize_named_params @@ -91,15 +116,27 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio return converted_named_params -def _quantize_param(name, weight, weight_block_size): +def _quantize_param( + name, + weight, + weight_block_size, + transform_ue8m0=True, + force_ue8m0_scale=False, +): assert name.endswith(".weight"), f"Expected weight parameter, got {name}" FP8_MIN = torch.finfo(torch.float8_e4m3fn).min FP8_MAX = torch.finfo(torch.float8_e4m3fn).max if weight_block_size is not None: - if should_deepgemm_weight_requant_ue8m0 and should_deepgemm_weight_requant_ue8m0( - weight_block_size=weight_block_size - ): + runtime_requires_ue8m0 = bool( + should_deepgemm_weight_requant_ue8m0 + and should_deepgemm_weight_requant_ue8m0(weight_block_size=weight_block_size) + ) + if force_ue8m0_scale or runtime_requires_ue8m0: qweight, scale = quant_weight_ue8m0(weight, weight_block_size=weight_block_size) + # Hopper keeps the power-of-two scales in the canonical FP32 block + # layout. Only the Blackwell DeepGEMM runtime consumes packed UE8M0. + if runtime_requires_ue8m0 and transform_ue8m0: + scale = transform_scale_ue8m0(scale, mn=qweight.shape[-2]) else: qweight, scale = blockwise_cast_to_fp8_triton(weight, weight_block_size) scale_name = name.replace(".weight", ".weight_scale_inv") diff --git a/vime/backends/megatron_utils/megatron_to_hf/qwen2.py b/vime/backends/megatron_utils/megatron_to_hf/qwen2.py index f7b72935c..bf14116eb 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/qwen2.py +++ b/vime/backends/megatron_utils/megatron_to_hf/qwen2.py @@ -61,6 +61,11 @@ def convert_qwen2_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] elif rest == "mlp.linear_fc1.layer_norm_weight": return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + # Local spec (--transformer-impl local) uses different layernorm names + elif rest == "input_layernorm.weight": + return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] + elif rest == "pre_mlp_layernorm.weight": + return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] # qk norm elif rest == "self_attention.q_layernorm.weight": diff --git a/vime/backends/megatron_utils/megatron_to_hf/qwen3_5.py b/vime/backends/megatron_utils/megatron_to_hf/qwen3_5.py index 2aabd86eb..eea892753 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/qwen3_5.py +++ b/vime/backends/megatron_utils/megatron_to_hf/qwen3_5.py @@ -37,6 +37,9 @@ def convert_qwen3_5_to_hf(args, name, param): Qwen3.5 uses model.language_model.layers prefix and has separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a for linear attention. """ + if name.startswith("module.module.language_model."): + name = "module.module." + name.removeprefix("module.module.language_model.") + # Handle MTP layers if "mtp.layers" in name: parts = name.split(".") diff --git a/vime/backends/megatron_utils/megatron_to_hf/qwen3_omni.py b/vime/backends/megatron_utils/megatron_to_hf/qwen3_omni.py new file mode 100644 index 000000000..cc6f5d660 --- /dev/null +++ b/vime/backends/megatron_utils/megatron_to_hf/qwen3_omni.py @@ -0,0 +1,14 @@ +from .qwen3moe import convert_qwen3moe_to_hf + + +def convert_qwen3_omni_to_hf(args, name, param): + prefixes = { + "module.module.audio_model.": "thinker.audio_tower.", + "module.module.vision_model.": "thinker.visual.", + } + for model_prefix, hf_prefix in prefixes.items(): + if name.startswith(model_prefix): + return [(hf_prefix + name.removeprefix(model_prefix), param)] + + name = name.replace("module.module.language_model.", "module.module.", 1) + return [("thinker." + hf_name, tensor) for hf_name, tensor in convert_qwen3moe_to_hf(args, name, param)] diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index 6b602fc1c..5f625e19f 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -5,6 +5,7 @@ import os from argparse import Namespace from collections.abc import Callable, Sequence +from contextlib import contextmanager, nullcontext from functools import partial from pathlib import Path @@ -27,14 +28,14 @@ from megatron.core.pipeline_parallel.utils import unwrap_model except ImportError: from megatron.core.utils import unwrap_model -from vime.utils import logging_utils +from vime.observability import logging_utils, train_metric_utils from vime.utils.memory_utils import clear_memory from .checkpoint import load_checkpoint, save_checkpoint -from .cp_utils import reduce_train_step_metrics from .data import DataIterator, get_batch -from .loss import loss_function +from .loss import ROLLOUT_TOP_P_TOKEN_KEYS, get_rollout_top_p_logprob_kwargs, loss_function from .model_provider import get_model_provider_func +from .stateless_adam import StatelessAdam logger = logging.getLogger(__name__) @@ -73,6 +74,12 @@ def wrapped_forward_step(*args, **kwargs): return wrapped_forward_step +def _with_rollout_top_p_token_keys(args: Namespace, keys: Sequence[str]) -> list[str]: + if args.rollout_top_p == 1.0: + return list(keys) + return [*keys, *ROLLOUT_TOP_P_TOKEN_KEYS] + + def _iter_critic_output_layers(model: Sequence[DDP]): for chunk_id, module in enumerate(unwrap_model(model)): output_layer = getattr(module, "output_layer", None) @@ -80,12 +87,45 @@ def _iter_critic_output_layers(model: Sequence[DDP]): yield chunk_id, output_layer +try: + from megatron.training.checkpointing import get_load_checkpoint_path_by_args +except ImportError: + + def get_load_checkpoint_path_by_args(args, load_arg="load"): + from megatron.training.checkpointing import ( + get_checkpoint_name, + get_checkpoint_tracker_filename, + isfile, + read_metadata, + ) + + """Get the checkpoint path based on the arguments.""" + load_dir = getattr(args, load_arg) + iteration, release = -1, False + tracker_filename = "because load directory is not defined" + if load_dir is not None: + tracker_filename = get_checkpoint_tracker_filename(load_dir) + if isfile(tracker_filename): + iteration, release = read_metadata(tracker_filename) + else: + load_dir, checkpoint_step = os.path.split(load_dir) + if checkpoint_step == "release" or checkpoint_step.startswith("iter_"): + release = checkpoint_step == "release" + if not release: + iteration = int(checkpoint_step.split("_")[1]) + + # Allow user to specify the loaded iteration. + if getattr(args, "ckpt_step", None): + iteration = args.ckpt_step + + return get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) + + def _critic_output_layer_needs_reinit(args: Namespace, model: Sequence[DDP], role: str) -> bool: if role != "critic" or args.load is None: return False from megatron.core.dist_checkpointing.serialization import load_tensors_metadata - from megatron.training.checkpointing import get_load_checkpoint_path_by_args checkpoint_path = Path(get_load_checkpoint_path_by_args(args)) if not (checkpoint_path / ".metadata").is_file(): @@ -128,9 +168,12 @@ def _critic_output_layer_needs_reinit(args: Namespace, model: Sequence[DDP], rol @torch.no_grad() -def _reinitialize_critic_output_layer(model: Sequence[DDP]) -> None: +def _reinitialize_critic_output_layer(args: Namespace, model: Sequence[DDP]) -> None: + init_method_std = getattr(args, "init_method_std", None) + if init_method_std is None: + init_method_std = 0.02 for _chunk_id, output_layer in _iter_critic_output_layers(model): - output_layer.weight.data.normal_(mean=0.0, std=0.02) + output_layer.weight.data.normal_(mean=0.0, std=init_method_std) if output_layer.bias is not None: output_layer.bias.data.zero_() @@ -191,10 +234,42 @@ def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) return opt_param_scheduler +def _noop_init_state_fn(*args, **kwargs) -> None: + return None + + +def _disable_distributed_optimizer_state_initialization(optimizer: MegatronOptimizer) -> None: + for megatron_optimizer in getattr(optimizer, "chained_optimizers", [optimizer]): + if megatron_optimizer.__class__.__name__ == "DistributedOptimizer": + megatron_optimizer.init_state_fn = _noop_init_state_fn + + +@contextmanager +def _patch_megatron_adam(adam_cls): + import megatron.core.optimizer as megatron_optimizer + import megatron.core.optimizer.distrib_optimizer as megatron_distrib_optimizer + + missing = object() + old_adam = megatron_optimizer.Adam + old_cpu_adam = getattr(megatron_optimizer, "CPUAdam", missing) + old_distrib_adam = megatron_distrib_optimizer.Adam + try: + megatron_optimizer.Adam = adam_cls + if old_cpu_adam is not missing: + megatron_optimizer.CPUAdam = adam_cls + megatron_distrib_optimizer.Adam = adam_cls + yield + finally: + megatron_optimizer.Adam = old_adam + if old_cpu_adam is not missing: + megatron_optimizer.CPUAdam = old_cpu_adam + megatron_distrib_optimizer.Adam = old_distrib_adam + + def setup_model_and_optimizer( args: Namespace, role: str = "actor", -) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]: +) -> tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None]: """Build model(s), wrap with DDP, and construct optimizer and scheduler. Args: @@ -207,7 +282,7 @@ def setup_model_and_optimizer( lr_mult (float): Global learning-rate multiplier for the optimizer. Returns: - tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler]: + tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None]: - List of model chunks wrapped by ``DDP``. - The constructed ``MegatronOptimizer`` instance. - The learning-rate/weight-decay scheduler tied to the optimizer. @@ -217,6 +292,10 @@ def setup_model_and_optimizer( model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder) + if args.num_rollout == 0: + args.no_load_optim = True + return model, None, None + # Optimizer kwargs = {} for f in dataclasses.fields(OptimizerConfig): @@ -225,11 +304,19 @@ def setup_model_and_optimizer( config = OptimizerConfig(**kwargs) config.timers = None - optimizer = get_megatron_optimizer( - config=config, - model_chunks=model, - use_gloo_process_groups=args.enable_gloo_process_groups, - ) + if args.use_stateless_adam: + assert config.optimizer == "adam", "Stateless Adam only supports --optimizer adam." + assert args.no_save_optim, "Stateless Adam does not save Adam moment states. Please set --no-save-optim." + + optimizer_context = _patch_megatron_adam(StatelessAdam) if args.use_stateless_adam else nullcontext() + with optimizer_context: + optimizer = get_megatron_optimizer( + config=config, + model_chunks=model, + use_gloo_process_groups=args.enable_gloo_process_groups, + ) + if args.use_stateless_adam: + _disable_distributed_optimizer_state_initialization(optimizer) opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer) return model, optimizer, opt_param_scheduler @@ -265,6 +352,7 @@ def forward_only( data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], store_prefix: str = "", + use_rollout_top_p_replay: bool = False, ) -> dict[str, list[torch.Tensor]]: """Run forward passes only and collect non-loss outputs (e.g., logprobs). @@ -284,6 +372,8 @@ def forward_only( data_iterator (Sequence[DataIterator]): Iterable(s) yielding batches for inference. num_microbatches (Sequence[int]): Number of microbatches per rollout step. store_prefix (str): Prefix to prepend to stored output keys. + use_rollout_top_p_replay (bool): Whether to pass rollout top-p token sets + to the post-forward log-prob callback when top-p rollout is enabled. Returns: dict[str, list[torch.Tensor]]: Aggregated outputs keyed by ``store_prefix + key``. @@ -294,6 +384,15 @@ def forward_only( iterator.reset() config = get_model_config(model[0]) + batch_keys = [ + "tokens", + "loss_masks", + "multimodal_train_inputs", + "total_lengths", + "response_lengths", + ] + if use_rollout_top_p_replay: + batch_keys = _with_rollout_top_p_token_keys(args, batch_keys) def forward_step( data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False @@ -315,16 +414,8 @@ def forward_step( # Get the batch. batch = get_batch( data_iterator, - [ - "tokens", - "loss_masks", - "multimodal_train_inputs", - "total_lengths", - "response_lengths", - "max_seq_lens", - ], + batch_keys, args.data_pad_size_multiplier, - args.qkv_format, args.allgather_cp, ) unconcat_tokens = batch["unconcat_tokens"] @@ -344,15 +435,17 @@ def forward_step( forward_kwargs.update(batch["multimodal_train_inputs"]) output_tensor = model(**forward_kwargs) - return output_tensor, partial( - f, - args=args, - unconcat_tokens=unconcat_tokens, - total_lengths=total_lengths, - response_lengths=response_lengths, - with_entropy=args.use_rollout_entropy, - max_seq_lens=batch.get("max_seq_lens", None), - ) + output_kwargs = { + "args": args, + "unconcat_tokens": unconcat_tokens, + "total_lengths": total_lengths, + "response_lengths": response_lengths, + "with_entropy": args.use_rollout_entropy, + } + if use_rollout_top_p_replay: + output_kwargs.update(get_rollout_top_p_logprob_kwargs(args, batch)) + + return output_tensor, partial(f, **output_kwargs) # Turn on evaluation mode which disables dropout. for model_module in model: @@ -486,25 +579,29 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p # Get the batch. batch = get_batch( data_iterator, - [ - "tokens", - "multimodal_train_inputs", - "packed_seq_params", - "total_lengths", - "response_lengths", - "loss_masks", - "log_probs", - "ref_log_probs", - "values", - "advantages", - "returns", - "rollout_log_probs", - "max_seq_lens", - "teacher_log_probs", - "rollout_mask_sums", - ], + _with_rollout_top_p_token_keys( + args, + [ + "tokens", + "multimodal_train_inputs", + "packed_seq_params", + "total_lengths", + "response_lengths", + "loss_masks", + "log_probs", + "ref_log_probs", + "values", + "advantages", + "returns", + "rollout_log_probs", + "teacher_log_probs", + "rollout_mask_sums", + # Only present when dumping train debug data; lets the loss + # snapshot each sample's log_probs keyed by rollout position. + *(["partition"] if args.save_debug_train_data is not None else []), + ], + ), args.data_pad_size_multiplier, - args.qkv_format, args.allgather_cp, ) @@ -533,17 +630,56 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "loss_mask": batch["full_loss_masks"], } + # MCore MambaModel.forward does not accept loss_mask; masking is + # applied by the loss function instead. + _m = model + while hasattr(_m, "module"): + _m = _m.module + # Signature does not change during training, so compute once and cache. + accepts_loss_mask = getattr(_m, "_vime_forward_accepts_loss_mask", None) + if accepts_loss_mask is None: + import inspect + + accepts_loss_mask = "loss_mask" in inspect.signature(_m.forward).parameters + _m._vime_forward_accepts_loss_mask = accepts_loss_mask + if not accepts_loss_mask: + forward_kwargs.pop("loss_mask", None) + if batch["multimodal_train_inputs"] is not None: forward_kwargs.update(batch["multimodal_train_inputs"]) if args.enable_mtp_training: forward_kwargs["mtp_kwargs"] = {"mtp_labels": batch["tokens"]} - output_tensor = model(**forward_kwargs) + dspark_outputs = None + dspark_config = None + if args.dspark_enabled: + from vime.backends.megatron_utils.dspark.hidden_capture import forward_with_dspark + + output_tensor, dspark_outputs, dspark_config = forward_with_dspark( + model, + forward_kwargs, + batch, + args.dspark_target_layer_ids, + ) + else: + output_tensor = model(**forward_kwargs) if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": os.environ["ROUTING_REPLAY_STAGE"] = old_stage + if dspark_outputs is not None: + from vime.backends.megatron_utils.dspark.loss import build_combined_loss_fn + + return output_tensor, build_combined_loss_fn( + loss_function, + args, + batch, + num_microbatches, + step_global_batch_size, + dspark_outputs, + dspark_config, + ) return output_tensor, partial(loss_function, args, batch, num_microbatches, step_global_batch_size) # Forward pass. @@ -594,7 +730,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p optimizer.zero_grad() if mpu.is_pipeline_last_stage(ignore_virtual=True): - loss_reduced = reduce_train_step_metrics( + loss_reduced = train_metric_utils.reduce_train_step_metrics( losses_reduced, calculate_per_token_loss=args.calculate_per_token_loss, step_global_batch_size=step_global_batch_size, @@ -682,7 +818,7 @@ def train( and mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1 ): - print("Reset optimizer states") + logger.info("Reset optimizer states") for chained_optimizer in optimizer.chained_optimizers: for group in chained_optimizer.optimizer.param_groups: if "step" in group: @@ -763,15 +899,15 @@ def train( torch.distributed.all_reduce(values, group=tracker.get("reduce_group")) if tracker.get("avg_group") is not None: torch.distributed.all_reduce(values, group=tracker["avg_group"], op=torch.distributed.ReduceOp.AVG) - # here we assume only one mtp layer - mtp_losses = (tracker["values"] * mtp_loss_scale).item() + # Multi-head MTP: tracker["values"] is [num_mtp_layers]; aggregate below. + mtp_losses = tracker["values"] * mtp_loss_scale MTPLossLoggingHelper.clean_loss_in_tracker() # CI check: verify MTP loss is within expected bounds if args.ci_test: from vime.backends.megatron_utils.ci_utils import check_mtp_loss - check_mtp_loss(mtp_losses) + check_mtp_loss(mtp_losses.sum().item()) # per train step log. if ( @@ -788,7 +924,9 @@ def train( } log_dict[f"train/{role_tag}grad_norm"] = grad_norm if args.enable_mtp_training: - log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses + for _i in range(mtp_losses.shape[0]): + log_dict[f"train/{role_tag}mtp_{_i + 1}_loss"] = mtp_losses[_i].item() + log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses.sum().item() for param_group_id, param_group in enumerate(optimizer.param_groups): log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group) @@ -799,7 +937,8 @@ def train( logging_utils.log(args, log_dict, step_key="train/step") if args.ci_test and "train/train_rollout_logprob_abs_diff" in log_dict: - assert log_dict["train/train_rollout_logprob_abs_diff"] <= 0.1, f"{log_dict=}" + threshold = args.ci_train_rollout_logprob_abs_diff_threshold + assert log_dict["train/train_rollout_logprob_abs_diff"] <= threshold, f"{threshold=} {log_dict=}" if args.ci_test and not args.ci_disable_kl_checker: if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict: @@ -874,60 +1013,9 @@ def save( enable_forward_pre_hook(model) -def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: - """Save Megatron model in HuggingFace format. - - Args: - model (Sequence[DDP]): Sequence of DDP-wrapped model chunks. - rollout_id (int): Rollout ID for path formatting. - """ - if args.megatron_to_hf_mode != "bridge": - try: - from vime.backends.megatron_utils.hf_checkpoint_saver import save_hf_model_direct - - save_hf_model_direct(args, rollout_id, model) - except Exception as e: - if ( - mpu.get_data_parallel_rank(with_context_parallel=True) == 0 - and mpu.get_tensor_model_parallel_rank() == 0 - ): - logger.error(f"Failed to save HuggingFace format: {e}") - return - - should_log = ( - mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 - ) - - try: - from megatron.bridge import AutoBridge - - from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_megatron_model - - path = Path(args.save_hf.format(rollout_id=rollout_id)) - - if should_log: - logger.info(f"Saving model in HuggingFace format to {path}") - - bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) - - path.mkdir(parents=True, exist_ok=True) - - with patch_megatron_model(model): - bridge.save_hf_pretrained( - model, - path=path, - ) - - if should_log: - logger.info(f"Successfully saved HuggingFace model to {path}") - except Exception as e: - if should_log: - logger.error(f"Failed to save HuggingFace format: {e}") - - def initialize_model_and_optimizer( args: Namespace, role: str = "actor" -) -> tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]: +) -> tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None, int]: """Initialize model(s), optimizer, scheduler, and load from checkpoint. Args: @@ -935,10 +1023,18 @@ def initialize_model_and_optimizer( role (str): Logical role of the model (e.g., "actor", "critic"). Returns: - tuple[list[DDP], MegatronOptimizer, OptimizerParamScheduler, int]: + tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None, int]: DDP-wrapped model chunks, optimizer, scheduler, and iteration index. """ + if torch.version.hip: + import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module + + from vime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync + + filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync + print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility") + model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role) model[0].role = role reinit_critic_output_layer = _critic_output_layer_needs_reinit(args, model, role) @@ -948,10 +1044,9 @@ def initialize_model_and_optimizer( optimizer, opt_param_scheduler, checkpointing_context={}, - skip_load_to_model_and_opt=False, ) if reinit_critic_output_layer: - _reinitialize_critic_output_layer(model) + _reinitialize_critic_output_layer(args, model) if (args.fp16 or args.bf16) and optimizer is not None: optimizer.reload_model_params() clear_memory() diff --git a/vime/backends/megatron_utils/model_provider.py b/vime/backends/megatron_utils/model_provider.py index 8c2b7a33f..5e8b70dfb 100644 --- a/vime/backends/megatron_utils/model_provider.py +++ b/vime/backends/megatron_utils/model_provider.py @@ -17,9 +17,40 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.arguments import core_transformer_config_from_args -from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config from vime.utils.misc import load_function +_INDEXER_DIRECT_SUBMODULE_NAMES = frozenset( + { + "wq_b", + "wk", + "k_norm", + "weights_proj", + "index_kpool_compress_ape", + "index_kpool_compress_gate", + } +) + + +def _is_indexer_parameter(name: str) -> bool: + """Return whether *name* belongs to a DSA indexer. + + The GLM plugin exposes indexer projections directly under + ``self_attention``. Megatron's upstream DSA implementation instead nests + them under ``self_attention.core_attention.indexer``. Keep the ownership + check structural so similarly named non-indexer projections stay trainable. + """ + + parts = name.split(".") + for index, part in enumerate(parts): + if part != "self_attention" or index + 1 >= len(parts): + continue + attention_parts = parts[index + 1 :] + if attention_parts[0] in _INDEXER_DIRECT_SUBMODULE_NAMES: + return True + if "indexer" in attention_parts[:-1]: + return True + return False + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): @@ -38,7 +69,10 @@ def __init__( if bias: self.bias.sequence_parallel = True - self.weight.data.normal_(mean=0.0, std=0.02) + init_method_std = getattr(config, "init_method_std", None) + if init_method_std is None: + init_method_std = 0.02 + self.weight.data.normal_(mean=0.0, std=init_method_std) if bias: self.bias.data.zero_() @@ -81,56 +115,6 @@ def wrapped_model_provider( return wrapped_model_provider - if args.megatron_to_hf_mode == "bridge": - from megatron.bridge import AutoBridge - - import vime_plugins.megatron_bridge # noqa: F401 # register custom bridges - - bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) - provider = bridge.to_megatron_provider(load_weights=False) - # TODO: we should not manually set this... - provider.tensor_model_parallel_size = args.tensor_model_parallel_size - provider.pipeline_model_parallel_size = args.pipeline_model_parallel_size - provider.expert_model_parallel_size = args.expert_model_parallel_size - provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size - provider.sequence_parallel = args.sequence_parallel - provider.context_parallel_size = args.context_parallel_size - provider.variable_seq_lengths = args.variable_seq_lengths - provider.gradient_accumulation_fusion = args.gradient_accumulation_fusion - if hasattr(args, "moe_token_dispatcher_type"): - provider.moe_token_dispatcher_type = args.moe_token_dispatcher_type - if getattr(args, "decoder_first_pipeline_num_layers", None) is not None: - provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers - if getattr(args, "decoder_last_pipeline_num_layers", None) is not None: - provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers - provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff - provider.freeze_language_model = False - provider.freeze_vision_model = False - provider.moe_permute_fusion = args.moe_permute_fusion - provider.recompute_granularity = args.recompute_granularity - provider.recompute_method = args.recompute_method - provider.recompute_num_layers = args.recompute_num_layers - for key, value in vars(args).items(): - if hasattr(provider, key): - continue - setattr(provider, key, value) - provider.finalize() - - if role == "critic": - _original_provide = provider.provide - - def _critic_provide(pre_process=True, post_process=True, vp_stage=None): - model = _original_provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) - if post_process: - model.output_layer = LinearForLastLayer( - input_size=model.config.hidden_size, output_size=1, config=model.config - ) - return model - - return _critic_provide - - return provider.provide - def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None) -> GPTModel: """Builds the model. @@ -148,6 +132,10 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage # Experimental loading arguments from yaml config: TransformerConfig = core_transformer_config_from_args(args) + # Older GLM Megatron forks consumed this flag from TransformerConfig. + # Preserve that contract for custom specs, while freeze_model_params() + # below provides the concrete implementation on current Megatron. + config.freeze_indexer = getattr(args, "freeze_indexer", False) if args.spec is not None: transformer_layer_spec = import_module(args.spec) @@ -190,6 +178,7 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage qk_layernorm=args.qk_layernorm, multi_latent_attention=args.multi_latent_attention, moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, + normalization=args.normalization, ) build_model_context = nullcontext @@ -246,6 +235,11 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage if post_process and role == "critic": model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config) + if getattr(args, "dspark_enabled", False) and role == "actor" and post_process: + from vime.backends.megatron_utils.dspark.modeling import attach_dspark_model + + attach_dspark_model(model, args, config) + return model return model_provider @@ -293,3 +287,22 @@ def freeze_model_params(model: GPTModel, args: argparse.Namespace): if re.search(pattern, name): param.requires_grad = False break + + if getattr(args, "freeze_indexer", False): + frozen_indexer_params = [] + has_self_attention_params = False + for name, param in model.named_parameters(): + has_self_attention_params |= "self_attention" in name.split(".") + if _is_indexer_parameter(name): + param.requires_grad = False + frozen_indexer_params.append(name) + + if has_self_attention_params and not frozen_indexer_params: + raise RuntimeError( + "--freeze-indexer was requested, but this model chunk has self-attention " + "parameters and no recognized DSA indexer parameters." + ) + + # Some pipeline stages may legitimately own no indexer weights, so an + # empty local tuple is not itself an error. + model._vime_frozen_indexer_param_names = tuple(frozen_indexer_params) diff --git a/vime/backends/megatron_utils/npu_attention_patch.py b/vime/backends/megatron_utils/npu_attention_patch.py index 184f5e717..6cfc23222 100644 --- a/vime/backends/megatron_utils/npu_attention_patch.py +++ b/vime/backends/megatron_utils/npu_attention_patch.py @@ -79,7 +79,3 @@ def npu_dot_product_attention_forward( from megatron.core.transformer.dot_product_attention import DotProductAttention DotProductAttention.forward = npu_dot_product_attention_forward -print( - "[NPU PATCH] DotProductAttention.forward replaced with npu_fusion_attention " "BEFORE megatron_adaptor import", - flush=True, -) diff --git a/vime/backends/megatron_utils/server/__init__.py b/vime/backends/megatron_utils/server/__init__.py new file mode 100644 index 000000000..379ce238b --- /dev/null +++ b/vime/backends/megatron_utils/server/__init__.py @@ -0,0 +1,13 @@ +"""Megatron teacher server utilities.""" + +from vime.backends.megatron_utils.server.arguments import ( + add_megatron_server_arguments, + configure_megatron_server_args, + validate_megatron_server_args, +) + +__all__ = [ + "add_megatron_server_arguments", + "configure_megatron_server_args", + "validate_megatron_server_args", +] diff --git a/vime/backends/megatron_utils/server/arguments.py b/vime/backends/megatron_utils/server/arguments.py new file mode 100644 index 000000000..125ea0642 --- /dev/null +++ b/vime/backends/megatron_utils/server/arguments.py @@ -0,0 +1,131 @@ +import argparse +import os +from typing import Any + + +def _env_bool(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.lower() in {"true", "1", "yes", "y", "on"} + + +def _non_negative_int(value: Any) -> int: + value = int(value) + if value < 0: + raise argparse.ArgumentTypeError("must be >= 0") + return value + + +def _positive_int(value: Any) -> int: + value = int(value) + if value <= 0: + raise argparse.ArgumentTypeError("must be > 0") + return value + + +def _positive_float(value: Any) -> float: + value = float(value) + if value <= 0: + raise argparse.ArgumentTypeError("must be > 0") + return value + + +def add_megatron_server_arguments(parser): + group = parser.add_argument_group("megatron server") + group.add_argument( + "--teacher-port", + type=_positive_int, + default=_positive_int(os.getenv("TEACHER_PORT", "7999")), + help="HTTP port for the Megatron teacher server.", + ) + group.add_argument( + "--teacher-warmup-port", + type=_positive_int, + default=_positive_int(os.getenv("TEACHER_WARMUP_PORT", "7999")), + help="Temporary HTTP port used by the Megatron teacher warmup server.", + ) + group.add_argument( + "--teacher-warmup-timeout-s", + type=_positive_int, + default=_positive_int(os.getenv("TEACHER_WARMUP_TIMEOUT_S", "3000")), + help="Timeout in seconds for the Megatron teacher warmup request.", + ) + group.add_argument( + "--teacher-sample-reduction-chunk-size", + type=_positive_int, + default=4096, + help="Row chunk size used while sampling from TP-sharded teacher logits.", + ) + group.add_argument( + "--teacher-label-reduction-chunk-size", + type=_positive_int, + default=4096, + help="Row chunk size used while gathering label-token logprobs from TP-sharded teacher logits.", + ) + group.add_argument( + "--megatron-server-max-length", + type=_non_negative_int, + default=_non_negative_int(os.getenv("MEGATRON_SERVER_MAX_LENGTH", "0")), + help="Reject teacher requests longer than this many tokens. Set 0 to disable.", + ) + group.add_argument( + "--megatron-server-update-timeout-s", + type=_positive_float, + default=_positive_float(os.getenv("MEGATRON_SERVER_UPDATE_TIMEOUT_S", "3600")), + help="Default timeout in seconds for /update_from_disk when the request does not override it.", + ) + group.add_argument( + "--megatron-server-warmup", + action=argparse.BooleanOptionalAction, + default=_env_bool("MEGATRON_SERVER_WARMUP", True), + help="Whether to run a local warmup request before serving traffic.", + ) + return parser + + +def configure_megatron_server_args(args): + args.debug_train_only = True + args.use_kl_loss = False + args.offload_train = False + args.use_dynamic_batch_size = False + args.use_wandb = False + args.kl_coef = 0 + args.use_opd = False + args.use_critic = False + args.keep_old_actor = False + args.no_load_optim = True + args.no_load_rng = True + # Keep this as a list (not str), otherwise freeze logic iterates over characters. + args.only_train_params_name_list = ["nothing_to_train"] + return args + + +def validate_megatron_server_args(args): + positive_fields = [ + "teacher_port", + "teacher_warmup_port", + "teacher_warmup_timeout_s", + "teacher_sample_reduction_chunk_size", + "teacher_label_reduction_chunk_size", + "megatron_server_update_timeout_s", + ] + non_negative_fields = [ + "megatron_server_max_length", + ] + + for name in positive_fields: + if getattr(args, name) <= 0: + raise ValueError(f"{name} must be > 0") + for name in non_negative_fields: + if getattr(args, name) < 0: + raise ValueError(f"{name} must be >= 0") + + if args.only_train_params_name_list != ["nothing_to_train"]: + raise ValueError("Megatron server must not train any parameters.") + if not args.debug_train_only: + raise ValueError("Megatron server requires debug_train_only=True.") + if args.use_kl_loss or args.use_opd or args.use_critic: + raise ValueError("Megatron server only supports teacher logprob prefill mode.") + + return args diff --git a/vime/backends/megatron_utils/server/logprob_utils.py b/vime/backends/megatron_utils/server/logprob_utils.py new file mode 100644 index 000000000..9a386d356 --- /dev/null +++ b/vime/backends/megatron_utils/server/logprob_utils.py @@ -0,0 +1,572 @@ +import logging +from functools import partial +from typing import Any + +import ray +import torch +import torch.distributed as dist +from megatron.core import mpu + +from vime.backends.megatron_utils.actor import MegatronTrainRayActor +from vime.backends.megatron_utils.cp_utils import all_gather_with_cp, get_logits_and_tokens_offset_with_cp +from vime.backends.megatron_utils.data import get_data_iterator +from vime.backends.megatron_utils.loss import get_log_probs_and_entropy, get_responses +from vime.backends.megatron_utils.model import forward_only +from vime.utils import accelerator + +logging.getLogger().setLevel(logging.WARNING) + + +@torch.no_grad() +def sample_from_vocab_parallel_logits_without_full_gather( + vocab_parallel_logits: torch.Tensor, + *, + sample_n: int, + tp_group: dist.ProcessGroup | None = None, + reduction_chunk_size: int = 4096, + global_vocab_size: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample token ids/log-probs from TP-sharded logits without full-vocab gather. + + This performs a two-stage sampling: + 1) Sample which TP rank owns each sample slot from per-rank probability mass. + 2) Each TP rank samples local vocab ids only for the slots assigned to it. + Final `(token_id, log_prob)` tensors are merged with TP all-reduce(max), so + no rank materializes full-vocab probabilities. + + Args: + vocab_parallel_logits: Shape `[num_tokens, vocab_per_tp]` logits shard. + sample_n: Number of samples per token position, with replacement. + tp_group: Tensor-parallel process group. If None, uses Megatron TP group. + reduction_chunk_size: Row chunk size for denominator computation. + global_vocab_size: Optional unpadded vocab size. If set, padded logits + (outside `[0, global_vocab_size)`) are excluded from sampling. + + Returns: + sampled_token_ids: `[num_tokens, sample_n]` global token ids. + sampled_log_probs: `[num_tokens, sample_n]` log-probabilities. + """ + if vocab_parallel_logits.dim() != 2: + raise ValueError(f"Expected 2D logits, got shape={tuple(vocab_parallel_logits.shape)}") + if sample_n < 0: + raise ValueError(f"sample_n must be >= 0, got {sample_n}") + if reduction_chunk_size <= 0: + raise ValueError(f"reduction_chunk_size must be > 0, got {reduction_chunk_size}") + + num_tokens, vocab_per_tp = vocab_parallel_logits.shape + device = vocab_parallel_logits.device + logits_dtype = vocab_parallel_logits.dtype + + if tp_group is None: + tp_group = mpu.get_tensor_model_parallel_group() + tp_rank = mpu.get_tensor_model_parallel_rank() + tp_size = mpu.get_tensor_model_parallel_world_size() + + reduction_dtype = torch.float32 if logits_dtype in (torch.float16, torch.bfloat16) else logits_dtype + + vocab_start = tp_rank * vocab_per_tp + valid_vocab_per_tp = vocab_per_tp + if global_vocab_size is not None: + if global_vocab_size < 0: + raise ValueError(f"global_vocab_size must be >= 0, got {global_vocab_size}") + valid_vocab_per_tp = max(min(global_vocab_size - vocab_start, vocab_per_tp), 0) + + if valid_vocab_per_tp == 0: + local_max = torch.full((num_tokens, 1), -torch.inf, dtype=logits_dtype, device=device) + elif valid_vocab_per_tp == vocab_per_tp: + local_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values + else: + local_max = vocab_parallel_logits[:, :valid_vocab_per_tp].max(dim=-1, keepdim=True).values + + global_max = local_max.clone() + if tp_size > 1: + dist.all_reduce(global_max, op=dist.ReduceOp.MAX, group=tp_group) + + local_exp_sums = torch.zeros((num_tokens, 1), dtype=reduction_dtype, device=device) + if valid_vocab_per_tp > 0: + for start in range(0, num_tokens, reduction_chunk_size): + end = min(start + reduction_chunk_size, num_tokens) + logits_part = vocab_parallel_logits[start:end, :valid_vocab_per_tp] + centered = logits_part.to(reduction_dtype) - global_max[start:end].to(reduction_dtype) + local_exp_sums[start:end] = centered.exp_().sum(dim=-1, keepdim=True) + + global_exp_sums = local_exp_sums.clone() + if tp_size > 1: + dist.all_reduce(global_exp_sums, op=dist.ReduceOp.SUM, group=tp_group) + denom = global_exp_sums.clamp_min(torch.finfo(global_exp_sums.dtype).tiny) + + local_tp_mass = (local_exp_sums / denom).squeeze(-1) + if tp_size > 1: + gathered_tp_masses = [torch.empty_like(local_tp_mass) for _ in range(tp_size)] + dist.all_gather(gathered_tp_masses, local_tp_mass.contiguous(), group=tp_group) + tp_masses = torch.stack(gathered_tp_masses, dim=-1) + else: + tp_masses = local_tp_mass.unsqueeze(-1) + + if tp_rank == 0: + tp_probs = tp_masses / tp_masses.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(tp_masses.dtype).tiny) + tp_assignments = torch.multinomial(tp_probs, num_samples=sample_n, replacement=True) + else: + tp_assignments = torch.empty((num_tokens, sample_n), dtype=torch.long, device=device) + if tp_size > 1: + dist.broadcast(tp_assignments, src=mpu.get_tensor_model_parallel_src_rank(), group=tp_group) + + owner_slots = tp_assignments.eq(tp_rank) + if valid_vocab_per_tp == 0 and owner_slots.any(): + raise RuntimeError("Received sample slots on a TP rank with zero valid vocab.") + + sampled_token_ids = torch.full((num_tokens, sample_n), -1, dtype=torch.long, device=device) + sampled_log_probs = torch.full((num_tokens, sample_n), -torch.inf, dtype=logits_dtype, device=device) + log_denom = denom.log() + + if valid_vocab_per_tp > 0: + local_slot_count = owner_slots.sum(dim=-1) + max_slot_count = int(local_slot_count.max().item()) + for slot_count in range(1, max_slot_count + 1): + row_ids = torch.nonzero(local_slot_count == slot_count, as_tuple=False).flatten() + if row_ids.numel() == 0: + continue + + row_logits = vocab_parallel_logits.index_select(0, row_ids)[:, :valid_vocab_per_tp] + row_max = global_max.index_select(0, row_ids).to(reduction_dtype) + local_row_max = row_logits.max(dim=-1, keepdim=True).values.to(reduction_dtype) + row_weights = (row_logits.to(reduction_dtype) - local_row_max).exp_() + local_ids = torch.multinomial(row_weights, num_samples=slot_count, replacement=True) + + slot_cols = torch.nonzero(owner_slots.index_select(0, row_ids), as_tuple=False)[:, 1].view(-1, slot_count) + + selected_logits = torch.gather(row_logits, dim=-1, index=local_ids) + selected_log_probs = ( + selected_logits.to(reduction_dtype) - row_max - log_denom.index_select(0, row_ids) + ).to(logits_dtype) + global_ids = local_ids + vocab_start + + sampled_token_ids[row_ids.unsqueeze(-1), slot_cols] = global_ids + sampled_log_probs[row_ids.unsqueeze(-1), slot_cols] = selected_log_probs + + if tp_size > 1: + dist.all_reduce(sampled_token_ids, op=dist.ReduceOp.MAX, group=tp_group) + dist.all_reduce(sampled_log_probs, op=dist.ReduceOp.MAX, group=tp_group) + + if (sampled_token_ids < 0).any(): + raise RuntimeError("Distributed TP sampling produced incomplete sample slots.") + + return sampled_token_ids, sampled_log_probs + + +@torch.no_grad() +def get_label_token_log_probs_from_vocab_parallel_logits( + vocab_parallel_logits: torch.Tensor, + label_token_ids: torch.Tensor, + *, + tp_group: dist.ProcessGroup | None = None, + reduction_chunk_size: int = 4096, +) -> torch.Tensor: + if vocab_parallel_logits.dim() != 2: + raise ValueError(f"Expected 2D logits, got shape={tuple(vocab_parallel_logits.shape)}") + if label_token_ids.dim() != 2: + raise ValueError(f"Expected 2D label_token_ids, got shape={tuple(label_token_ids.shape)}") + if vocab_parallel_logits.size(0) != label_token_ids.size(0): + raise ValueError( + "label_token_ids must align with response positions: " + f"{vocab_parallel_logits.size(0)} vs {label_token_ids.size(0)}" + ) + if reduction_chunk_size <= 0: + raise ValueError(f"reduction_chunk_size must be > 0, got {reduction_chunk_size}") + + num_tokens, partition_vocab_size = vocab_parallel_logits.shape + num_labels = label_token_ids.size(1) + if num_labels == 0: + return torch.empty((num_tokens, 0), dtype=vocab_parallel_logits.dtype, device=vocab_parallel_logits.device) + + device = vocab_parallel_logits.device + logits_dtype = vocab_parallel_logits.dtype + reduction_dtype = torch.float32 if logits_dtype in (torch.float16, torch.bfloat16) else logits_dtype + + if tp_group is None: + tp_group = mpu.get_tensor_model_parallel_group() + tp_rank = mpu.get_tensor_model_parallel_rank() + tp_size = mpu.get_tensor_model_parallel_world_size() + + from megatron.core.tensor_parallel.utils import VocabUtility # type: ignore + + vocab_start_index, vocab_end_index = VocabUtility.vocab_range_from_per_partition_vocab_size( + partition_vocab_size, tp_rank, tp_size + ) + + local_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values + global_max = local_max.clone() + if tp_size > 1: + dist.all_reduce(global_max, op=dist.ReduceOp.MAX, group=tp_group) + + local_exp_sums = torch.zeros((num_tokens, 1), dtype=reduction_dtype, device=device) + for start in range(0, num_tokens, reduction_chunk_size): + end = min(start + reduction_chunk_size, num_tokens) + logits_part = vocab_parallel_logits[start:end] + centered = logits_part.to(reduction_dtype) - global_max[start:end].to(reduction_dtype) + local_exp_sums[start:end] = centered.exp_().sum(dim=-1, keepdim=True) + + global_exp_sums = local_exp_sums.clone() + if tp_size > 1: + dist.all_reduce(global_exp_sums, op=dist.ReduceOp.SUM, group=tp_group) + log_denom = global_exp_sums.clamp_min(torch.finfo(global_exp_sums.dtype).tiny).log() + + local_mask = (label_token_ids >= vocab_start_index) & (label_token_ids < vocab_end_index) + local_label_ids = (label_token_ids - vocab_start_index).masked_fill(~local_mask, 0) + local_label_ids = local_label_ids.clamp_(0, max(partition_vocab_size - 1, 0)) + + local_selected_logits = torch.gather(vocab_parallel_logits, dim=-1, index=local_label_ids) + local_selected_logits = local_selected_logits.masked_fill(~local_mask, 0.0).to(reduction_dtype) + if tp_size > 1: + dist.all_reduce(local_selected_logits, op=dist.ReduceOp.SUM, group=tp_group) + + return (local_selected_logits - global_max.to(reduction_dtype) - log_denom).to(logits_dtype) + + +def _to_accelerator_tensors(values, dtype: torch.dtype) -> list[torch.Tensor]: + return [torch.as_tensor(value, dtype=dtype, device=accelerator.current_device()) for value in values] + + +def _prepare_rollout_data(rollout_data_ref): + rollout_data = ray.get(rollout_data_ref[0].inner) + + rollout_data["tokens"] = _to_accelerator_tensors(rollout_data["tokens"], torch.long) + rollout_data["loss_masks"] = _to_accelerator_tensors(rollout_data["loss_masks"], torch.int) + if rollout_data.get("label_token_ids") is not None: + rollout_data["label_token_ids"] = _to_accelerator_tensors(rollout_data["label_token_ids"], torch.long) + for idx, tensor in enumerate(rollout_data["label_token_ids"]): + if tensor.dim() == 1 and tensor.numel() == 0: + rollout_data["label_token_ids"][idx] = tensor.reshape(0, 0) + + micro_batch_size = len(rollout_data["tokens"]) + rollout_data["micro_batch_indices"] = [list(range(micro_batch_size))] + rollout_data["num_microbatches"] = [1] + rollout_data["global_batch_sizes"] = [micro_batch_size] + + return rollout_data + + +def _merge_tensors_with_cp( + tensors: list[torch.Tensor] | None, + total_lengths: list[int], + response_lengths: list[int], + args, +) -> list[torch.Tensor] | None: + if not tensors: + return tensors + + cp_size = mpu.get_context_parallel_world_size() + if cp_size == 1: + return tensors + + merged = [] + for tensor, total_length, response_length in zip(tensors, total_lengths, response_lengths, strict=False): + merged.append(all_gather_with_cp(tensor, total_length, response_length)) + return merged + + +def _slice_response_rows_for_current_cp_rank( + rows: torch.Tensor, + sample_idx: int, + *, + logits_local_len: int, + args, + total_lengths: list[int], + response_lengths: list[int], +) -> torch.Tensor: + cp_size = mpu.get_context_parallel_world_size() + if cp_size == 1: + return rows + + total_length = total_lengths[sample_idx] + response_length = response_lengths[sample_idx] + prompt_length = total_length - response_length + + if getattr(args, "allgather_cp", False): + seq_start = sum(total_lengths[:sample_idx]) + chunk_start = mpu.get_context_parallel_rank() * logits_local_len + chunk_end = chunk_start + logits_local_len + logit_global_start = seq_start + prompt_length - 1 + logit_global_end = seq_start + total_length - 1 + start = max(logit_global_start, chunk_start) + end = min(logit_global_end, chunk_end) + if end <= start: + return rows[:0] + return rows[start - logit_global_start : end - logit_global_start] + + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length) + rows_0 = rows[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + rows_1 = rows[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + return torch.cat([rows_0, rows_1], dim=0) + + +def _merge_allgather_cp_tensors( + outputs: dict[str, list[torch.Tensor]], + keys: tuple[str, ...], + *, + logits_local_len: int, + total_lengths: list[int], + response_lengths: list[int], +) -> None: + if mpu.get_context_parallel_world_size() == 1: + return + + cp_rank = mpu.get_context_parallel_rank() + cp_group = mpu.get_context_parallel_group() + chunk_start = cp_rank * logits_local_len + chunk_end = chunk_start + logits_local_len + + for key in keys: + values = outputs.get(key) + if values is None: + continue + + full_values = [] + seq_start = 0 + for value, total_length, response_length in zip(values, total_lengths, response_lengths, strict=False): + prompt_length = total_length - response_length + logit_global_start = seq_start + prompt_length - 1 + logit_global_end = seq_start + total_length - 1 + start = max(logit_global_start, chunk_start) + end = min(logit_global_end, chunk_end) + + if end <= start: + full_value = value.new_zeros((response_length, *value.shape[1:])) + else: + expected_len = end - start + if value.size(0) != expected_len: + raise ValueError(f"{key} length mismatch: got {value.size(0)}, expected {expected_len}") + response_start = start - logit_global_start + response_end = end - logit_global_start + left = value.new_zeros((response_start, *value.shape[1:])) + right = value.new_zeros((response_length - response_end, *value.shape[1:])) + full_value = torch.cat([left, value, right], dim=0) + + full_values.append(full_value) + seq_start += total_length + + gathered = dist.nn.all_reduce(torch.cat(full_values, dim=0), group=cp_group) + outputs[key] = list(gathered.split(response_lengths, dim=0)) + + +def _get_log_probs_and_optional_samples( + logits: torch.Tensor, + *, + args, + unconcat_tokens: list[torch.Tensor], + total_lengths: list[int], + response_lengths: list[int], + with_entropy: bool = False, + non_loss_data: bool = True, + sample_n: int = 0, + label_token_ids: list[torch.Tensor] | None = None, +) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]: + _, outputs = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=with_entropy, + non_loss_data=non_loss_data, + ) + logits_local_len = logits.size(1) + + if label_token_ids is not None: + if len(label_token_ids) != len(unconcat_tokens): + raise ValueError(f"label_token_ids batch size mismatch: {len(label_token_ids)} vs {len(unconcat_tokens)}") + label_token_log_probs = [] + label_reduction_chunk_size = args.teacher_label_reduction_chunk_size + for sample_idx, ((logits_chunk, _), sample_label_token_ids) in enumerate( + zip( + get_responses( + logits, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + ), + label_token_ids, + strict=True, + ) + ): + local_label_token_ids = _slice_response_rows_for_current_cp_rank( + sample_label_token_ids, + sample_idx, + logits_local_len=logits_local_len, + args=args, + total_lengths=total_lengths, + response_lengths=response_lengths, + ) + label_token_log_probs.append( + get_label_token_log_probs_from_vocab_parallel_logits( + logits_chunk, + local_label_token_ids, + reduction_chunk_size=label_reduction_chunk_size, + ) + ) + outputs["label_token_log_probs"] = label_token_log_probs + + if sample_n > 0: + sampled_token_ids = [] + sampled_log_probs = [] + reduction_chunk_size = args.teacher_sample_reduction_chunk_size + + for logits_chunk, _ in get_responses( + logits, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + ): + if logits_chunk.size(0) == 0: + sampled_token_ids.append(torch.empty((0, sample_n), dtype=torch.long, device=logits_chunk.device)) + sampled_log_probs.append( + torch.empty((0, sample_n), dtype=logits_chunk.dtype, device=logits_chunk.device) + ) + continue + + sampled_ids, sampled_logp = sample_from_vocab_parallel_logits_without_full_gather( + logits_chunk, + sample_n=sample_n, + reduction_chunk_size=reduction_chunk_size, + ) + + sampled_token_ids.append(sampled_ids) + sampled_log_probs.append(sampled_logp) + + outputs["sampled_token_ids"] = sampled_token_ids + outputs["sampled_log_probs"] = sampled_log_probs + + if getattr(args, "allgather_cp", False): + if "sampled_token_ids" in outputs: + outputs["sampled_token_ids"] = [x.to(torch.float32) for x in outputs["sampled_token_ids"]] + _merge_allgather_cp_tensors( + outputs, + ("label_token_log_probs", "sampled_token_ids", "sampled_log_probs"), + logits_local_len=logits_local_len, + total_lengths=total_lengths, + response_lengths=response_lengths, + ) + return torch.empty((0,), device=logits.device), outputs + + +class TeacherLogpRayActor(MegatronTrainRayActor): + """A Megatron actor subclass that exposes a log-prob computation RPC.""" + + def get_parallel_infos(self) -> dict[str, int]: + return { + "dp_rank": mpu.get_data_parallel_rank(), + "pp_rank": mpu.get_pipeline_model_parallel_rank(), + "tp_rank": mpu.get_tensor_model_parallel_rank(), + "cp_rank": mpu.get_context_parallel_rank(), + "dp_size": mpu.get_data_parallel_world_size(), + "cp_size": mpu.get_context_parallel_world_size(), + "tp_size": mpu.get_tensor_model_parallel_world_size(), + "pp_size": mpu.get_pipeline_model_parallel_world_size(), + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + logging.getLogger().setLevel(logging.WARNING) + + def compute_logp(self, rollout_data_ref) -> dict[str, Any]: + rollout_data = _prepare_rollout_data(rollout_data_ref) + sample_ns = rollout_data.get("sample_ns", [0]) + sample_n = int(sample_ns[0]) if sample_ns else 0 + label_token_ids = rollout_data.get("label_token_ids") + if sample_n < 0: + raise ValueError(f"sample_n must be >= 0, got {sample_n}") + + data_iterator = get_data_iterator(rollout_data) + num_microbatches = rollout_data["num_microbatches"] + + if sample_n > 0 or label_token_ids is not None: + forward_outputs = forward_only( + partial( + _get_log_probs_and_optional_samples, + sample_n=sample_n, + label_token_ids=label_token_ids, + ), + self.args, + self.model, + data_iterator, + num_microbatches, + store_prefix="", + ) + else: + forward_outputs = self.compute_log_prob( + data_iterator, + num_microbatches, + store_prefix="", + ) + + log_probs = forward_outputs.get("log_probs", None) + sampled_token_ids = forward_outputs.get("sampled_token_ids", None) + sampled_log_probs = forward_outputs.get("sampled_log_probs", None) + label_token_log_probs = forward_outputs.get("label_token_log_probs", None) + if mpu.is_pipeline_last_stage(): + log_probs = _merge_tensors_with_cp( + log_probs, + rollout_data["total_lengths"], + rollout_data["response_lengths"], + self.args, + ) + if sampled_log_probs is not None: + if not self.args.allgather_cp: + sampled_log_probs = _merge_tensors_with_cp( + sampled_log_probs, + rollout_data["total_lengths"], + rollout_data["response_lengths"], + self.args, + ) + if sampled_token_ids is not None: + sampled_token_ids = [x.to(torch.float32) for x in sampled_token_ids] + if not self.args.allgather_cp: + sampled_token_ids = _merge_tensors_with_cp( + sampled_token_ids, + rollout_data["total_lengths"], + rollout_data["response_lengths"], + self.args, + ) + if label_token_log_probs is not None: + if not self.args.allgather_cp: + label_token_log_probs = _merge_tensors_with_cp( + label_token_log_probs, + rollout_data["total_lengths"], + rollout_data["response_lengths"], + self.args, + ) + if mpu.get_context_parallel_rank() == 0 and mpu.get_tensor_model_parallel_rank() == 0: + log_prob = log_probs[0].tolist() + sampled_log_prob = sampled_log_probs[0].tolist() if sampled_log_probs else None + sampled_token_id = sampled_token_ids[0].round().to(torch.long).tolist() if sampled_token_ids else None + label_token_log_prob = label_token_log_probs[0].tolist() if label_token_log_probs else None + else: + log_prob = None + sampled_log_prob = None + sampled_token_id = None + label_token_log_prob = None + else: + log_prob = None + sampled_log_prob = None + sampled_token_id = None + label_token_log_prob = None + + return { + "log_prob": log_prob, + "sampled_log_probs": sampled_log_prob, + "sampled_token_ids": sampled_token_id, + "label_token_log_probs": label_token_log_prob, + "dp_rank": mpu.get_data_parallel_rank(with_context_parallel=False), + "pp_rank": mpu.get_pipeline_model_parallel_rank(), + "cp_rank": mpu.get_context_parallel_rank(), + "tp_rank": mpu.get_tensor_model_parallel_rank(), + } + + def update_from_disk(self, model_path: str) -> dict[str, Any]: + self.load_other_checkpoint("actor", model_path) + return { + "rank": self.args.rank, + "model_path": model_path, + } diff --git a/vime/backends/megatron_utils/server/megatron_server.py b/vime/backends/megatron_utils/server/megatron_server.py new file mode 100644 index 000000000..e28564071 --- /dev/null +++ b/vime/backends/megatron_utils/server/megatron_server.py @@ -0,0 +1,770 @@ +import asyncio +import sys +import threading +import time +from collections import defaultdict, deque +from datetime import datetime +from pathlib import Path +from typing import Any + +import httpx +import ray +from aiohttp import web + + +def _find_project_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "vime").is_dir() and (parent / "setup.py").is_file(): + return parent + return Path(__file__).resolve().parents[4] + + +PROJECT_ROOT = _find_project_root() +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from vime.agent.aiohttp_threaded import FilteredAccessLogger, run_app_in_thread # noqa: E402 +from vime.backends.megatron_utils.server.arguments import ( # noqa: E402 + add_megatron_server_arguments, + configure_megatron_server_args, + validate_megatron_server_args, +) +from vime.utils.misc import Box # noqa: E402 + + +def _parse_sample_n(sample_n: Any) -> int: + if sample_n is None: + return 0 + if isinstance(sample_n, bool): + raise ValueError("sample_n must be an integer >= 0") + try: + sample_n = int(sample_n) + except (TypeError, ValueError) as exc: + raise ValueError("sample_n must be an integer >= 0") from exc + if sample_n < 0: + raise ValueError("sample_n must be >= 0") + return sample_n + + +def _normalize_label_token_ids(label_token_ids: Any, expected_length: int) -> list[list[int]] | None: + if label_token_ids is None: + return None + if not isinstance(label_token_ids, list): + raise ValueError("label_token_ids must be a 2D list with shape [len(input_ids) - 1, num_label_tokens]") + if len(label_token_ids) != expected_length: + raise ValueError("label_token_ids length must match len(input_ids) - 1") + + normalized = [] + width = None + for row in label_token_ids: + if not isinstance(row, list): + raise ValueError("label_token_ids must be a 2D list with shape [len(input_ids) - 1, num_label_tokens]") + if width is None: + width = len(row) + elif len(row) != width: + raise ValueError("label_token_ids rows must have the same length") + try: + normalized.append([int(token_id) for token_id in row]) + except (TypeError, ValueError) as exc: + raise ValueError("label_token_ids must contain integers") from exc + + return normalized + + +def _get_max_request_length(args) -> int: + return args.megatron_server_max_length + + +def _get_update_timeout_s(payload: dict[str, Any], args) -> float: + return float(payload.get("timeout_s") or args.megatron_server_update_timeout_s) + + +@ray.remote +class SampleManager: + """Minimal rollout-manager surface plus request queue for teacher-server mode.""" + + def __init__(self, args, pg): + self.args = args + self.pg = pg + self.train_parallel_config = None + + self._pending_requests = deque() + self._inflight = defaultdict(deque) # worker_id -> deque of batch_info + self._results = {} + self._next_request_id = 0 + self._canceled = set() + self.total_finished_reqs = 0 + self.total_finished_tokens = 0 + + def set_train_parallel_config(self, config: dict): + self.train_parallel_config = config + + def submit( + self, + input_ids, + request_id=None, + response_length=None, + loss_mask=None, + metadata=None, + sample_n=0, + label_token_ids=None, + ): + if len(input_ids) == 0: + raise ValueError("input_ids is empty") + sample_n = _parse_sample_n(sample_n) + + if response_length is None: + response_length = max(len(input_ids) - 1, 0) + + if loss_mask is None: + loss_mask = [1] * response_length + + if response_length > len(input_ids) - 1: + raise ValueError("response_length error") + if len(loss_mask) != response_length: + raise ValueError("loss_mask length error") + + if request_id is None: + rid = f"req_{self._next_request_id}" + self._next_request_id += 1 + else: + rid = request_id + + self._pending_requests.append( + { + "request_id": rid, + "tokens": input_ids, + "response_length": response_length, + "loss_mask": loss_mask, + "metadata": metadata, + "sample_n": sample_n, + "label_token_ids": label_token_ids, + } + ) + return rid + + def _stats(self): + total_inflight = sum(len(q) for q in self._inflight.values()) + return { + "queue_size": len(self._pending_requests), + "inflight_size": total_inflight, + } + + def get_stats(self): + return self._stats() + + def get_global_stats(self): + return { + **self._stats(), + "total_finished_reqs": self.total_finished_reqs, + "total_finished_tokens": self.total_finished_tokens, + } + + def get_loads(self): + pending_tokens = sum(len(req["tokens"]) for req in self._pending_requests) + pending_response_tokens = sum(req["response_length"] for req in self._pending_requests) + running_by_worker = {worker_id: len(queue) for worker_id, queue in self._inflight.items()} + running_tokens = sum(sum(info["response_lengths"]) for queue in self._inflight.values() for info in queue) + return { + **self.get_global_stats(), + "pending_tokens": pending_tokens, + "pending_response_tokens": pending_response_tokens, + "running_tokens": running_tokens, + "running_by_worker": running_by_worker, + } + + def cancel_request(self, request_id: str) -> None: + self._canceled.add(request_id) + self._results.pop(request_id, None) + if self._pending_requests: + self._pending_requests = deque( + req for req in self._pending_requests if req.get("request_id") != request_id + ) + + def _pop_next_request(self): + while self._pending_requests: + req = self._pending_requests.popleft() + rid = req.get("request_id") + if rid in self._canceled: + continue + return req + return None + + def get_input_data(self, worker_id): + if self.train_parallel_config is None or not self._pending_requests: + return None + + req = self._pop_next_request() + if req is None: + return None + + tokens = req["tokens"] + response_length = req["response_length"] + loss_mask = req["loss_mask"] + sample_n = req.get("sample_n", 0) + label_token_ids = req.get("label_token_ids") + + request_ids = [req["request_id"]] + is_dummy = [False] + token_lens = [len(tokens)] + response_lengths = [response_length] + + rollout_data = { + "tokens": [tokens], + "response_lengths": [response_length], + "loss_masks": [loss_mask], + "sample_indices": [0], + "total_lengths": [len(tokens)], + "sample_ns": [sample_n], + } + if label_token_ids is not None: + rollout_data["label_token_ids"] = [label_token_ids] + data_refs = [Box(ray.put(rollout_data))] + + self._inflight[worker_id].append( + { + "request_ids": request_ids, + "is_dummy": is_dummy, + "token_lens": token_lens, + "response_lengths": response_lengths, + "sample_ns": [sample_n], + } + ) + + return data_refs + + def load(self, rollout_id=None): + pass + + def save_log_probs(self, worker_id, outputs): + if not self._inflight[worker_id]: + raise KeyError(f"No inflight info for worker {worker_id}") + + info = self._inflight[worker_id].popleft() + + request_ids = info["request_ids"] + is_dummy = info["is_dummy"] + response_lengths = info["response_lengths"] + + assert len(outputs) == len(request_ids) + + for i, rid in enumerate(request_ids): + if is_dummy[i]: + continue + + self.total_finished_reqs += 1 + self.total_finished_tokens += response_lengths[i] + + if rid in self._canceled: + continue + + output_item = outputs[i] + if isinstance(output_item, dict): + log_probs = output_item.get("log_probs") + sampled_token_ids = output_item.get("sampled_token_ids") + sampled_log_probs = output_item.get("sampled_log_probs") + label_token_log_probs = output_item.get("label_token_log_probs") + else: + # Keep compatibility with the legacy output format. + log_probs = output_item + sampled_token_ids = None + sampled_log_probs = None + label_token_log_probs = None + + self._results[rid] = { + "log_probs": log_probs, + "sampled_token_ids": sampled_token_ids, + "sampled_log_probs": sampled_log_probs, + "label_token_log_probs": label_token_log_probs, + } + return [rid for rid in request_ids if rid is not None] + + def get_result(self, request_id): + return self._results.pop(request_id, None) + + +def _merge_log_probs(logp_parts: list[dict[str, Any]]) -> list[dict[str, Any]]: + logp_parts.sort(key=lambda x: x.get("dp_rank")) + merged = [] + for part in logp_parts: + has_log_prob = part.get("log_prob") is not None + has_sampled = part.get("sampled_log_probs") is not None or part.get("sampled_token_ids") is not None + has_label = part.get("label_token_log_probs") is not None + if not has_log_prob and not has_sampled and not has_label: + continue + merged.append( + { + "log_probs": part.get("log_prob"), + "sampled_log_probs": part.get("sampled_log_probs"), + "sampled_token_ids": part.get("sampled_token_ids"), + "label_token_log_probs": part.get("label_token_log_probs"), + } + ) + return merged + + +def _run_stats_printer(sample_manager, interval=1.0): + """像 VLLM 一样每隔一段时间打印系统状态""" + last_time = time.time() + last_reqs = 0 + last_tokens = 0 + + print(f"Stats printer started (interval={interval}s)...", flush=True) + + while True: + time.sleep(interval) + try: + # 获取全局统计 + stats = ray.get(sample_manager.get_global_stats.remote()) + + now = time.time() + delta_t = now - last_time + if delta_t <= 0: + continue + + # 计算增量 + curr_reqs = stats["total_finished_reqs"] + curr_tokens = stats["total_finished_tokens"] + + delta_reqs = curr_reqs - last_reqs + delta_tokens = curr_tokens - last_tokens + + rps = delta_reqs / delta_t + tps = delta_tokens / delta_t + + # 格式化输出 + current_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + if stats["queue_size"] > 0 or stats["inflight_size"] > 0 or rps > 0 or tps > 0: + print( + f"[{current_time_str}] " + f"Queue: {stats['queue_size']} | " + f"Running: {stats['inflight_size']} | " + f"RPS: {rps:.2f} | " + f"TPS: {tps:.2f} tokens/s", + flush=True, + ) + + last_time = now + last_reqs = curr_reqs + last_tokens = curr_tokens + + except Exception as e: + print(f"Stats printer error: {e}", flush=True) + + +async def _ray_get(ref): + return await asyncio.to_thread(ray.get, ref) + + +async def _read_json_payload(request: web.Request) -> dict[str, Any]: + if not request.can_read_body: + return {} + try: + payload = await request.json() + except ValueError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _json_error(message: str, status: int) -> web.Response: + return web.json_response({"error": message}, status=status) + + +def _crop_sequence(value, expected_len: int): + if value is not None and len(value) > expected_len: + return value[:expected_len] + return value + + +def _build_generate_response(request_id: str, result: dict[str, Any], expected_len: int) -> dict[str, Any]: + response = { + "request_id": request_id, + "log_probs": _crop_sequence(result.get("log_probs"), expected_len), + } + for key in ("label_token_log_probs", "sampled_token_ids", "sampled_log_probs"): + value = result.get(key) + if value is not None: + response[key] = _crop_sequence(value, expected_len) + return response + + +async def _wait_until_idle(sample_manager, timeout_s: float) -> dict[str, Any]: + deadline = time.time() + timeout_s + while True: + loads = await _ray_get(sample_manager.get_loads.remote()) + if loads["queue_size"] == 0 and loads["inflight_size"] == 0: + return loads + if time.time() >= deadline: + raise TimeoutError(f"timed out waiting for queued/inflight requests to finish: {loads}") + await asyncio.sleep(0.1) + + +def _get_update_model_path(payload: dict[str, Any]) -> str | None: + model_path = payload.get("model_path") or payload.get("path") or payload.get("load") + return str(model_path) if model_path else None + + +def _jsonable(value: Any) -> Any: + """Best-effort conversion of an arbitrary value to something JSON-serializable.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_jsonable(v) for v in value] + return str(value) + + +def _args_to_dict(args) -> dict[str, Any]: + return {key: _jsonable(val) for key, val in sorted(vars(args).items())} + + +def _build_http_app(sample_manager, args, update_from_disk_fn=None): + app = web.Application(client_max_size=64 * 1024 * 1024) + update_state = { + "in_progress": False, + "updating_model_path": None, + "update_future": None, + } + update_lock = asyncio.Lock() + + async def detect(_request: web.Request) -> web.Response: + return web.json_response({"server_type": "megatron_server"}) + + async def healthz(_request: web.Request) -> web.Response: + return web.json_response({"ok": True}) + + async def info(_request: web.Request) -> web.Response: + return web.json_response({"args": _args_to_dict(args)}) + + async def get_loads(_request: web.Request) -> web.Response: + return web.json_response(await _ray_get(sample_manager.get_loads.remote())) + + async def update_from_disk(request: web.Request) -> web.Response: + if update_from_disk_fn is None: + return _json_error("update_from_disk is not available during warmup", 503) + + payload = await _read_json_payload(request) + model_path = _get_update_model_path(payload) + if model_path is None: + return _json_error("missing model_path", 400) + + async with update_lock: + if getattr(args, "load", None) == model_path: + return web.json_response({"ok": True, "model_path": model_path, "skipped": True}) + + if update_state["in_progress"]: + if update_state["updating_model_path"] == model_path and update_state["update_future"] is not None: + update_future = update_state["update_future"] + coalesced = True + else: + updating_model_path = update_state["updating_model_path"] + return _json_error(f"update_from_disk is already in progress for {updating_model_path}", 409) + else: + update_future = asyncio.get_running_loop().create_future() + update_state["in_progress"] = True + update_state["updating_model_path"] = model_path + update_state["update_future"] = update_future + coalesced = False + + if coalesced: + result = await asyncio.shield(update_future) + if result.get("ok") is True: + result = dict(result) + result["coalesced"] = True + return web.json_response(result) + return _json_error(result.get("error", "update_from_disk failed"), int(result.get("status", 500))) + + timeout_s = _get_update_timeout_s(payload, args) + result = None + error = None + try: + before_loads = await _wait_until_idle(sample_manager, timeout_s) + update_result = await asyncio.to_thread(update_from_disk_fn, model_path) + after_loads = await _ray_get(sample_manager.get_loads.remote()) + except TimeoutError as e: + error = {"ok": False, "status": 503, "error": str(e)} + except Exception as e: + error = {"ok": False, "status": 500, "error": f"update_from_disk failed: {e}"} + finally: + if error is None: + result = { + "ok": True, + "model_path": model_path, + "before_loads": before_loads, + "after_loads": after_loads, + "update_result": update_result, + } + # Reflect the freshly loaded checkpoint in /info. The actors restore + # their own args after loading, so only the server-side copy needs to be + # kept in sync here. + args.load = model_path + args.ref_load = model_path + + async with update_lock: + if result is not None: + update_future.set_result(result) + else: + update_future.set_result(error) + update_state["in_progress"] = False + update_state["updating_model_path"] = None + update_state["update_future"] = None + + if result is not None: + return web.json_response(result) + return _json_error(error["error"], error["status"]) + + async def generate(request: web.Request) -> web.Response: + if update_state["in_progress"]: + return _json_error("server is updating from disk", 503) + + payload = await _read_json_payload(request) + if "input_ids" not in payload: + return _json_error("missing input_ids", 400) + + input_ids = payload["input_ids"] + original_input_len = len(input_ids) + max_request_length = _get_max_request_length(args) + if max_request_length > 0 and original_input_len > max_request_length: + return web.json_response( + { + "error": ( + f"input_ids length {original_input_len} exceeds configured maximum {max_request_length}" + ), + "input_length": original_input_len, + "max_length": max_request_length, + }, + status=413, + ) + + try: + sample_n = _parse_sample_n(payload.get("sample_n", 0)) + valid_response_length = max(original_input_len - 1, 0) + label_token_ids = _normalize_label_token_ids(payload.get("label_token_ids"), valid_response_length) + except ValueError as e: + return _json_error(str(e), 400) + + try: + stats = await _ray_get(sample_manager.get_stats.remote()) + current_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + print( + f"[{current_time_str}] Received Request | " + f"Queue: {stats['queue_size']} | " + f"Running: {stats['inflight_size']} | " + f"Input Len: {original_input_len} | " + f"Max Len: {max_request_length or 'disabled'} | " + f"sample_n: {sample_n} | " + f"label_width: {len(label_token_ids[0]) if label_token_ids else 0}", + flush=True, + ) + except Exception as e: + print(f"stats failed: {e}", flush=True) + + request_id = None + try: + if update_state["in_progress"]: + return _json_error("server is updating from disk", 503) + + response_length = valid_response_length + loss_mask = [1] * response_length + request_id = await _ray_get( + sample_manager.submit.remote( + input_ids, + response_length=response_length, + loss_mask=loss_mask, + sample_n=sample_n, + label_token_ids=label_token_ids, + ) + ) + + while True: + result = await _ray_get(sample_manager.get_result.remote(request_id)) + if result is not None: + break + await asyncio.sleep(0.05) + except asyncio.CancelledError: + if request_id is not None: + print(f"cancel request {request_id} because client disconnected", flush=True) + await _ray_get(sample_manager.cancel_request.remote(request_id)) + raise + except Exception as e: + prefix = "submit failed" if request_id is None else "result failed" + return _json_error(f"{prefix}: {e}", 500) + + expected_log_probs_len = max(original_input_len - 1, 0) + return web.json_response(_build_generate_response(request_id, result, expected_log_probs_len)) + + app.router.add_get("/detect", detect) + app.router.add_get("/healthz", healthz) + app.router.add_get("/info", info) + app.router.add_get("/get_loads", get_loads) + app.router.add_post("/update_weights_from_disk", update_from_disk) + app.router.add_post("/generate", generate) + return app + + +def _start_http_server(sample_manager, args, update_from_disk_fn): + app = _build_http_app(sample_manager, args, update_from_disk_fn=update_from_disk_fn) + handle = run_app_in_thread( + app, + host="0.0.0.0", + port=args.teacher_port, + thread_name="teacher-http", + runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, + ) + + stats_thread = threading.Thread( + target=_run_stats_printer, + args=(sample_manager, 5.0), + name="stats-printer", + daemon=True, + ) + stats_thread.start() + return handle + + +def _run_warmup_via_private_http(sample_manager, args): + warmup_timeout_s = args.teacher_warmup_timeout_s + warmup_host = "127.0.0.1" + warmup_port = args.teacher_warmup_port + warmup_tokens = [101, 102, 103, 104] + + print( + "start private warmup server: " + f"http://{warmup_host}:{warmup_port}, timeout_s={warmup_timeout_s}, input_len={len(warmup_tokens)}", + flush=True, + ) + + app = _build_http_app(sample_manager, args) + warmup_handle = run_app_in_thread( + app, + host=warmup_host, + port=warmup_port, + thread_name="teacher-warmup-http", + runner_kwargs={"handler_cancellation": True, "access_log_class": FilteredAccessLogger}, + ) + + try: + req_start = time.time() + timeout = httpx.Timeout(warmup_timeout_s + 60, connect=10.0) + with httpx.Client(timeout=timeout, trust_env=False) as client: + response = client.post( + f"http://{warmup_host}:{warmup_handle.port}/generate", + json={ + "input_ids": warmup_tokens, + "sample_n": 0, + "timeout_s": warmup_timeout_s, + }, + ) + if response.status_code != 200: + raise RuntimeError(f"warmup request failed with http {response.status_code}, body={response.text[:500]}") + + print( + f"private warmup request finished in {time.time() - req_start:.2f}s", + flush=True, + ) + finally: + warmup_handle.stop() + print("private warmup server stopped", flush=True) + + +@ray.remote(num_cpus=0) +def run_megatron_dp_models_loop_worker(dp_rank, pp_size, sample_manager, actor_models): + worker_id = f"rank_{dp_rank}" + print(f"Start Async Pipeline Loop for {worker_id}", flush=True) + + MAX_INFLIGHT_BATCHES = pp_size + 1 + futures_queue = deque() + + while True: + # 1. Submission Stage + if len(futures_queue) < MAX_INFLIGHT_BATCHES: + rollout_data_ref = ray.get(sample_manager.get_input_data.remote(worker_id)) + + if rollout_data_ref is not None: + logp_parts_refs = [actor.compute_logp.remote(rollout_data_ref) for actor in actor_models] + futures_queue.append(logp_parts_refs) + else: + if not futures_queue: + time.sleep(0.02) + + # 2. Collection Stage + if futures_queue: + oldest_refs = futures_queue[0] + should_block = len(futures_queue) >= MAX_INFLIGHT_BATCHES + + _, remaining_refs = ray.wait( + oldest_refs, num_returns=len(oldest_refs), timeout=None if should_block else 0 + ) + + if len(remaining_refs) == 0: + logp_parts = ray.get(oldest_refs) + futures_queue.popleft() + + merged_log_probs = _merge_log_probs(logp_parts) + sample_manager.save_log_probs.remote(worker_id, merged_log_probs) + + +def _build_update_from_disk_fn(actor_model): + def update_from_disk(model_path: str): + refs = [actor.update_from_disk.remote(model_path) for actor in actor_model._actor_handlers] + results = ray.get(refs) + return { + "num_ranks": len(results), + "model_path": model_path, + "results": results, + } + + return update_from_disk + + +def launch(args): + from vime.backends.megatron_utils.server.logprob_utils import TeacherLogpRayActor + from vime.ray.placement_group import create_placement_groups, create_training_models + + configure_megatron_server_args(args) + validate_megatron_server_args(args) + pgs = create_placement_groups(args) + + sample_manager = SampleManager.options( + num_cpus=1, + num_gpus=0, + ).remote(args, pgs["rollout"]) + + print("initializing training models...", flush=True) + actor_model, _ = create_training_models(args, pgs, sample_manager, actor_cls=TeacherLogpRayActor) + parallel_infos = ray.get([actor.get_parallel_infos.remote() for actor in actor_model._actor_handlers]) + all_dp_ranks = set(info["dp_rank"] for info in parallel_infos) + futures = [] + pp_size = parallel_infos[0]["pp_size"] + + for dp_rank in all_dp_ranks: + models_in_same_dp_rank = [ + actor + for info, actor in zip(parallel_infos, actor_model._actor_handlers, strict=True) + if info["dp_rank"] == dp_rank + ] + + # 注意：这里已经应用了之前讨论的异步 Worker + fut = run_megatron_dp_models_loop_worker.remote(dp_rank, pp_size, sample_manager, models_in_same_dp_rank) + futures.append(fut) + + if args.megatron_server_warmup: + _run_warmup_via_private_http(sample_manager, args) + else: + print("warmup disabled", flush=True) + print("training models ready and warmup done", flush=True) + _start_http_server(sample_manager, args, update_from_disk_fn=_build_update_from_disk_fn(actor_model)) + + ray.get(futures) + + +def main(): + from vime.utils.arguments import parse_args + + args = parse_args(add_custom_arguments=add_megatron_server_arguments) + launch(args) + + +if __name__ == "__main__": + main() diff --git a/vime/backends/megatron_utils/stateless_adam.py b/vime/backends/megatron_utils/stateless_adam.py new file mode 100644 index 000000000..4c49ab92b --- /dev/null +++ b/vime/backends/megatron_utils/stateless_adam.py @@ -0,0 +1,109 @@ +import math +from typing import Any + +import torch + + +class StatelessAdam(torch.optim.Optimizer): + """Adam/AdamW update for the special case where moments are reset every step. + + This optimizer intentionally does not keep ``exp_avg`` or ``exp_avg_sq``. + Its parameter update matches Adam with zero first and second moments at the + start of every optimizer step, which is the behavior produced by resetting + Adam state before each one-step rollout. + """ + + def __init__( + self, + params, + lr: float = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + amsgrad: bool = False, + *, + bias_correction: bool = True, + adam_w_mode: bool = True, + maximize: bool = False, + use_decoupled_grad: bool = False, + master_weights: bool = False, + set_grad_none: bool | None = None, + **_: Any, + ) -> None: + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if eps < 0.0: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta1 value: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta2 value: {betas[1]}") + if amsgrad: + raise NotImplementedError("StatelessAdam does not support amsgrad.") + if master_weights: + raise NotImplementedError("StatelessAdam relies on Megatron/HybridDeviceOptimizer for master weights.") + + defaults = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + amsgrad=amsgrad, + bias_correction=bias_correction, + adam_w_mode=adam_w_mode, + maximize=maximize, + use_decoupled_grad=use_decoupled_grad, + set_grad_none=True if set_grad_none is None else set_grad_none, + ) + super().__init__(params, defaults) + for group in self.param_groups: + group.setdefault("step", 0) + + def load_state_dict(self, state_dict) -> None: + state_dict = dict(state_dict) + state_dict["state"] = {} + super().load_state_dict(state_dict) + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + beta1, beta2 = group["betas"] + eps = group["eps"] + weight_decay = group["weight_decay"] + adam_w_mode = group.get("adam_w_mode", True) + bias_correction = group.get("bias_correction", True) + maximize = group.get("maximize", False) + use_decoupled_grad = group.get("use_decoupled_grad", False) + + group["step"] = 1 + + if bias_correction: + numerator_scale = 1.0 + denominator_scale = 1.0 + else: + numerator_scale = 1.0 - beta1 + denominator_scale = math.sqrt(1.0 - beta2) + + for param in group["params"]: + grad = getattr(param, "decoupled_grad", None) if use_decoupled_grad else param.grad + if grad is None: + continue + if grad.is_sparse: + raise RuntimeError("StatelessAdam does not support sparse gradients.") + + grad_for_update = grad.neg() if maximize else grad + if weight_decay != 0 and adam_w_mode: + param.mul_(1.0 - lr * weight_decay) + elif weight_decay != 0: + grad_for_update = grad_for_update.add(param, alpha=weight_decay) + + denom = grad_for_update.abs().mul(denominator_scale).add_(eps) + param.addcdiv_(grad_for_update, denom, value=-lr * numerator_scale) + + return loss diff --git a/vime/backends/megatron_utils/update_weight/__init__.py b/vime/backends/megatron_utils/update_weight/__init__.py index e69de29bb..dc40d630e 100644 --- a/vime/backends/megatron_utils/update_weight/__init__.py +++ b/vime/backends/megatron_utils/update_weight/__init__.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +import torch + + +def create_weight_updater( + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, Any] | None, +): + """Select and construct the weight updater for the configured transport.""" + update_weight_mode = args.update_weight_mode + update_weight_transport = args.update_weight_transport + + if update_weight_mode == "delta": + # Delta sync is disk-transport only: each engine's /pull_weights applies the published + # deltas into a host-local checkpoint on every host it spans, and the engines reload + # via vanilla update_weights_from_disk. + assert not args.colocate, "--update-weight-mode=delta is not supported with --colocate" + assert update_weight_transport == "disk", "--update-weight-mode=delta requires --update-weight-transport=disk" + from .update_weight_from_disk_delta import UpdateWeightFromDiskDelta + + update_weight_cls = UpdateWeightFromDiskDelta + elif update_weight_transport == "disk": + from .update_weight_from_disk import UpdateWeightFromDisk + + update_weight_cls = UpdateWeightFromDisk + elif args.colocate: + from .update_weight_from_tensor import UpdateWeightFromTensor + + update_weight_cls = UpdateWeightFromTensor + else: + assert update_weight_mode == "full" + assert ( + update_weight_transport == "nccl" + ), f"unsupported weight sync mode/transport: {update_weight_mode!r}/{update_weight_transport!r}" + from .update_weight_from_distributed import UpdateWeightFromDistributed + + update_weight_cls = UpdateWeightFromDistributed + + updater = update_weight_cls( + args, + model, + weights_getter, + model_name=model_name, + quantization_config=quantization_config, + ) + updater.weight_version = getattr(args, "update_weight_start_version", 0) + return updater diff --git a/vime/backends/megatron_utils/update_weight/common.py b/vime/backends/megatron_utils/update_weight/common.py index 8cbb025df..548f148c9 100644 --- a/vime/backends/megatron_utils/update_weight/common.py +++ b/vime/backends/megatron_utils/update_weight/common.py @@ -1,21 +1,24 @@ import inspect import re +import socket from argparse import Namespace -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import Any import torch import torch.distributed as dist from megatron.core import mpu from megatron.core.transformer.transformer_layer import get_transformer_layer_offset -from vime.backends.megatron_utils.misc_utils import strip_param_name_prefix -from vime.utils.common import is_npu +from vime.platforms import current_platform +from vime.utils.distributed_utils import get_gloo_group from vime.utils.types import ParamInfo def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: """ - All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data. + All-gather TP-sharded param to full tensor. expert_bias→param, + non-TP/duplicated/TP-size-1→param.data. Uses expert-TP for ".experts.", else regular-TP. linear_fc1 rechunked (GLU), linear_fc2 dim fix. """ if "expert_bias" in name: @@ -27,9 +30,15 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: if ".experts." in name: tp_size = mpu.get_expert_tensor_parallel_world_size() - tp_group = mpu.get_expert_tensor_parallel_group() else: tp_size = mpu.get_tensor_model_parallel_world_size() + + if tp_size == 1: + return param.data + + if ".experts." in name: + tp_group = mpu.get_expert_tensor_parallel_group() + else: tp_group = mpu.get_tensor_model_parallel_group() param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] @@ -43,8 +52,7 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: if "linear_fc1.weight" in name or "linear_fc1.bias" in name: param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] - if is_npu(): - partition_dim = 0 + partition_dim = current_platform().megatron.adjust_tp_partition_dim(name, partition_dim) # this is bug in megatron's grouped moe. if "linear_fc2.weight" in name: if partition_dim == 0: @@ -58,7 +66,8 @@ def all_gather_params_async( ) -> list[torch.Tensor]: """ Parallel TP all-gather for multiple params. Loop 1: for each TP param, allocate buffers + - dist.all_gather(async_op=True) on expert-TP/regular-TP group (skip expert_bias/non-TP/duplicated). + dist.all_gather(async_op=True) on expert-TP/regular-TP group + (skip expert_bias/non-TP/duplicated/TP-size-1). Loop 2: wait all NCCL handles (enables overlap). Loop 3: concat partitions + apply GLU rechunk/MoE dim fix. """ # Phase 1: Start all async all_gather operations @@ -69,17 +78,22 @@ def all_gather_params_async( # Prepare async all_gather if "expert_bias" in info.name: gather_tasks.append((info, param, None, None, None)) - handles.append(None) elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated": gather_tasks.append((info, param.data, None, None, None)) - handles.append(None) else: # Start async all_gather if ".experts." in info.name: tp_size = mpu.get_expert_tensor_parallel_world_size() - tp_group = mpu.get_expert_tensor_parallel_group() else: tp_size = mpu.get_tensor_model_parallel_world_size() + + if tp_size == 1: + gather_tasks.append((info, param.data, None, None, None)) + continue + + if ".experts." in info.name: + tp_group = mpu.get_expert_tensor_parallel_group() + else: tp_group = mpu.get_tensor_model_parallel_group() param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] @@ -90,8 +104,7 @@ def all_gather_params_async( # Phase 2: Wait for ALL async operations to complete at once # This ensures maximum parallelism by not blocking on individual operations for handle in handles: - if handle is not None: - handle.wait() + handle.wait() # Phase 3: Process all results after all communications are done gathered_params = [] @@ -107,8 +120,7 @@ def all_gather_params_async( if "linear_fc1.weight" in info.name or "linear_fc1.bias" in info.name: param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] - if is_npu(): - partition_dim = 0 + partition_dim = current_platform().megatron.adjust_tp_partition_dim(info.name, partition_dim) # this is bug in megatron's grouped moe. if "linear_fc2.weight" in info.name: if partition_dim == 0: @@ -120,51 +132,7 @@ def all_gather_params_async( return gathered_params -def named_params_and_buffers( - args: Namespace, - model: Sequence[torch.nn.Module], - convert_to_global_name: bool = True, - translate_gpu_to_cpu: bool = False, -) -> Iterator[tuple[str, torch.Tensor]]: - if convert_to_global_name: - ans = _named_params_and_buffers_global(args, model) - else: - ans = _named_params_and_buffers_vanilla(model) - - if translate_gpu_to_cpu: - ans = ((name, _maybe_get_cpu_backup(tensor)) for name, tensor in ans) - - return ans - - -def _maybe_get_cpu_backup(x: torch.Tensor): - from torch_memory_saver import torch_memory_saver - - if (cpu_tensor := torch_memory_saver.get_cpu_backup(x, zero_copy=True)) is not None: - return cpu_tensor - - return x - - -def _named_params_and_buffers_vanilla(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]: - for vp_stage, model_module in enumerate(model): - - def _compute_fqn(name, vp_stage=vp_stage): - return f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}" - - for name, param in model_module.named_parameters(): - yield _compute_fqn(name), param - - for name, buffer in model_module.named_buffers(): - # TODO shall we handle (almost) all buffers like Megatron Bridge - if "expert_bias" not in name: - continue - yield _compute_fqn(name), buffer - - -def _named_params_and_buffers_global( - args: Namespace, model: Sequence[torch.nn.Module] -) -> Iterator[tuple[str, torch.Tensor]]: +def named_params_and_buffers(args: Namespace, model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]: """ Yield (global_name, param/buffer) with consistent names across PP/EP. Adjusts indices for virtual PP + EP offsets. Handles decoder.layers, mtp.layers (Multi-Token Prediction), expert_bias. @@ -186,12 +154,13 @@ def _named_params_and_buffers_global( # for model without ddp wrap if not name.startswith("module.module."): name = "module." + name + prefix = "module.module.language_model." if ".language_model." in name else "module.module." - decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" + decoder_layers_pattern = r"module\.module\.(?:language_model\.)?decoder\.layers\.(\d+)\.(.+)" match = re.match(decoder_layers_pattern, name) if not match: # MTP (Multi-Token Prediction) layers for speculative decoding - mtp_layers_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)" + mtp_layers_pattern = r"module\.module\.(?:language_model\.)?mtp\.layers\.(\d+)\.(.+)" match = re.match(mtp_layers_pattern, name) if not match: yield name, param @@ -207,7 +176,7 @@ def _named_params_and_buffers_global( rest, param_type, expert_idx = match.groups() expert_idx = int(expert_idx) + expert_offset - yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.{param_type}{expert_idx}", param + yield f"{prefix}mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.{param_type}{expert_idx}", param continue layer_idx, rest = match.groups() @@ -219,24 +188,127 @@ def _named_params_and_buffers_global( if match: rest, param_type, expert_idx = match.groups() expert_idx = int(expert_idx) + expert_offset - yield f"module.module.decoder.layers.{layer_idx}.mlp.experts.{rest}.{param_type}{expert_idx}", param + yield f"{prefix}decoder.layers.{layer_idx}.mlp.experts.{rest}.{param_type}{expert_idx}", param else: - yield f"module.module.decoder.layers.{layer_idx}.{rest}", param + yield f"{prefix}decoder.layers.{layer_idx}.{rest}", param # treat expert bias as normal parameters for name, buffer in model_module.named_buffers(): - # TODO shall we handle (almost) all buffers like Megatron Bridge + # TODO shall we handle (almost) all buffers if "expert_bias" not in name: continue # for model without ddp wrap if not name.startswith("module.module."): name = "module." + name + prefix = "module.module.language_model." if ".language_model." in name else "module.module." - decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" + decoder_layers_pattern = r"module\.module\.(?:language_model\.)?decoder\.layers\.(\d+)\.(.+)" match = re.match(decoder_layers_pattern, name) if not match: yield name, buffer else: layer_idx, rest = match.groups() layer_idx = int(layer_idx) + layer_offset - yield f"module.module.decoder.layers.{layer_idx}.{rest}", buffer + yield f"{prefix}decoder.layers.{layer_idx}.{rest}", buffer + + +class HfWeightSource: + def __init__( + self, + iterator, + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + draft_weights_getter: Callable[[], Sequence[tuple[str, torch.Tensor]]] | None = None, + ) -> None: + self.iterator = iterator + self.weights_getter = weights_getter + self.draft_weights_getter = draft_weights_getter + self.draft = False + self._metadata = {} + + def metadata(self): + if self.draft not in self._metadata: + from vllm.distributed.weight_transfer.base import ParamMeta + + self._metadata[self.draft] = [ParamMeta(name, tensor.dtype, tuple(tensor.shape)) for name, tensor in self] + return self._metadata[self.draft] + + def __iter__(self): + if self.draft: + if self.draft_weights_getter is None: + raise RuntimeError("Draft weight update requested without a draft weight source") + yield from self.draft_weights_getter() + return + for chunk in self.iterator.get_hf_weight_chunks(self.weights_getter()): + yield from chunk + + +class VimeRayWeightSyncClient: + def __init__( + self, + engines: Sequence[Any], + version_getter: Callable[[], int], + engine_gpu_counts: Sequence[int] | None = None, + ) -> None: + self.engines = list(engines) + self.version_getter = version_getter + self.engine_gpu_counts = engine_gpu_counts + self.draft = False + + def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: + import ray + + refs = [] + rank_offset = 1 + for index, engine in enumerate(self.engines): + engine_info = dict(init_info) + if self.engine_gpu_counts is not None: + engine_info["rank_offset"] = rank_offset + rank_offset += self.engine_gpu_counts[index] + refs.append(engine.init_weight_transfer_engine.remote({"init_info": engine_info})) + ray.get(refs) + + def start_weight_update(self) -> None: + import ray + + method = "start_draft_weight_update" if self.draft else "start_weight_update" + ray.get([getattr(engine, method).remote() for engine in self.engines]) + + def update_weights(self, update_info: dict[str, Any] | list[dict[str, Any]]) -> None: + import ray + + ray.get([engine.update_weights.remote(update_info) for engine in self.engines]) + + def finish_weight_update(self, weight_version: str | None = None) -> None: + import ray + + version = str(self.version_getter()) if weight_version is None else str(weight_version) + ray.get([engine.finish_weight_update.remote(weight_version=version) for engine in self.engines]) + + +def create_nccl_trainer( + client: VimeRayWeightSyncClient, + source: HfWeightSource, + engine_gpu_counts: Sequence[int], +): + import ray + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + rendezvous = [None] + if dist.get_rank() == 0: + with socket.socket() as sock: + sock.bind(("", 0)) + rendezvous[0] = (ray._private.services.get_node_ip_address(), sock.getsockname()[1]) + dist.broadcast_object_list(rendezvous, src=0, group=get_gloo_group()) + master_address, master_port = rendezvous[0] + return WeightTransferTrainerFactory.trainer_init( + current_platform().weight_transfer.trainer_init_info( + colocate=False, + master_address=master_address, + master_port=master_port, + world_size=sum(engine_gpu_counts) + 1, + rank=dist.get_rank(), + packed_num_buffers=1, + ), + client=client, + source=source, + ) diff --git a/vime/backends/megatron_utils/update_weight/expert_routing.py b/vime/backends/megatron_utils/update_weight/expert_routing.py new file mode 100644 index 000000000..b77419bb1 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/expert_routing.py @@ -0,0 +1,413 @@ +import logging +import re +from argparse import Namespace +from collections import defaultdict +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Any + +import torch.distributed as dist + +from vime.utils.distributed_utils import get_gloo_group +from vime.utils.types import ParamInfo + +__all__ = ["configure_expert_routing"] + + +logger = logging.getLogger(__name__) + +_ROUTED_EXPERT = re.compile(r"module\.module\.decoder\.layers\.(\d+)\.mlp\.experts\.linear_fc([12])\.weight(\d+)") + + +@dataclass(frozen=True) +class _ExpertParam: + info: ParamInfo + layer: int + expert: int + target_ranks: tuple[int, ...] + + +@dataclass(frozen=True) +class _ExpertTransfer: + source_rank: int + target_ranks: tuple[int, ...] + params: tuple[_ExpertParam, ...] + + +_ExpertTransferBatch = tuple[_ExpertTransfer, ...] +_ExpertTransferGroup = tuple[_ExpertTransferBatch, ...] + + +@dataclass(frozen=True) +class _VLLMMoeTopology: + tp_size: int + pp_size: int + pcp_size: int + dp_size: int + enable_expert_parallel: bool + ep_size: int + + +def _config_value( + parallel_config: Mapping[str, Any] | None, + key: str, + default: Any, +) -> Any: + if parallel_config is None: + return default + return parallel_config.get(key, parallel_config.get(key.replace("_", "-"), default)) + + +def _get_vllm_moe_topology( + args: Namespace, + engine_gpu_count: int, + parallel_config: Mapping[str, Any] | None = None, +) -> _VLLMMoeTopology: + pp_size = int(_config_value(parallel_config, "pp_size", getattr(args, "vllm_pp_size", 1)) or 1) + pcp_size = int( + _config_value( + parallel_config, + "pcp_size", + getattr(args, "vllm_prefill_context_parallel_size", 1), + ) + or 1 + ) + dp_size = int(_config_value(parallel_config, "dp_size", getattr(args, "vllm_dp_size", 1)) or 1) + parallel_divisor = pp_size * pcp_size * dp_size + if engine_gpu_count % parallel_divisor: + raise ValueError( + f"VLLM engine GPU count {engine_gpu_count} is not divisible by PP*PCP*DP " + f"({pp_size}*{pcp_size}*{dp_size})" + ) + default_tp_size = engine_gpu_count // parallel_divisor + tp_size = int(_config_value(parallel_config, "tp_size", default_tp_size) or default_tp_size) + if tp_size * parallel_divisor != engine_gpu_count: + raise ValueError( + f"VLLM engine GPU count {engine_gpu_count} does not match TP*PP*PCP*DP " + f"({tp_size}*{pp_size}*{pcp_size}*{dp_size})" + ) + enable_expert_parallel = bool( + _config_value( + parallel_config, + "enable_expert_parallel", + getattr(args, "vllm_enable_expert_parallel", False), + ) + ) + ep_size = tp_size * pcp_size * dp_size if enable_expert_parallel else 1 + + return _VLLMMoeTopology( + tp_size=tp_size, + pp_size=pp_size, + pcp_size=pcp_size, + dp_size=dp_size, + enable_expert_parallel=enable_expert_parallel, + ep_size=ep_size, + ) + + +def _vllm_topology_signature(topology: _VLLMMoeTopology) -> tuple[int, int, int, int, bool, int]: + return ( + topology.tp_size, + topology.pp_size, + topology.pcp_size, + topology.dp_size, + topology.enable_expert_parallel, + topology.ep_size, + ) + + +def _get_homogeneous_vllm_moe_topology( + args: Namespace, + engine_gpu_counts: Sequence[int], + engine_parallel_configs: Sequence[Mapping[str, Any]] | None, +) -> _VLLMMoeTopology: + if engine_parallel_configs is None: + return _get_vllm_moe_topology(args, engine_gpu_count=engine_gpu_counts[0]) + if len(engine_parallel_configs) != len(engine_gpu_counts): + raise ValueError( + f"VLLM engine parallel config count {len(engine_parallel_configs)} " + f"!= engine count {len(engine_gpu_counts)}" + ) + + topologies = [ + _get_vllm_moe_topology(args, engine_gpu_count=gpu_count, parallel_config=parallel_config) + for gpu_count, parallel_config in zip(engine_gpu_counts, engine_parallel_configs, strict=True) + ] + signatures = {_vllm_topology_signature(topology) for topology in topologies} + if len(signatures) != 1: + raise ValueError(f"VLLM engines have heterogeneous parallel topology: {sorted(signatures)}") + return topologies[0] + + +def _can_route_experts( + args: Namespace, + vllm_moe_topology: _VLLMMoeTopology, + engine_gpu_counts: Sequence[int], +) -> bool: + from megatron.core import mpu + + eplb_config = getattr(args, "vllm_eplb_config", None) + if isinstance(eplb_config, Mapping): + num_redundant_experts = eplb_config.get("num_redundant_experts", 0) + else: + num_redundant_experts = getattr(eplb_config, "num_redundant_experts", 0) + + return ( + vllm_moe_topology.pp_size == 1 + and vllm_moe_topology.enable_expert_parallel + and vllm_moe_topology.ep_size > 1 + and not getattr(args, "vllm_enable_eplb", False) + and num_redundant_experts == 0 + and getattr(args, "vllm_expert_placement_strategy", "linear") == "linear" + and not getattr(args, "vllm_enable_elastic_ep", False) + and mpu.get_expert_tensor_parallel_world_size() == 1 + and _vllm_moe_tp_is_one(engine_gpu_counts, vllm_moe_topology) + ) + + +def _vllm_moe_tp_is_one( + engine_gpu_counts: Sequence[int], + topology: _VLLMMoeTopology, +) -> bool: + """Return whether each VLLM engine has no tensor parallelism inside experts.""" + if topology.pp_size != 1: + return False + expected_size = topology.pp_size * topology.ep_size + return all(gpu_count == expected_size for gpu_count in engine_gpu_counts) + + +def _get_expert_target_ranks( + engine_gpu_counts: Sequence[int], + engine_gpu_offsets: Sequence[int], + *, + ep_size: int, + world_size: int, +) -> tuple[tuple[int, ...], ...]: + """Map each EP shard to the corresponding colocated rank.""" + expected_size = ep_size + targets = [[] for _ in range(ep_size)] + for gpu_count, gpu_offset in zip(engine_gpu_counts, engine_gpu_offsets, strict=True): + if gpu_count != expected_size: + raise ValueError(f"VLLM MoE TP must be 1, got engine_size={gpu_count}, EP={ep_size}") + if gpu_offset < 0 or gpu_offset + gpu_count > world_size: + raise ValueError("VLLM engine is outside the Megatron world") + for ep_rank in range(ep_size): + targets[ep_rank].append(gpu_offset + ep_rank) + return tuple(tuple(ranks) for ranks in targets) + + +def _build_expert_params( + infos: Sequence[ParamInfo], + target_ranks: Sequence[Sequence[int]], + *, + num_experts: int, +) -> list[_ExpertParam]: + ep_size = len(target_ranks) + if num_experts % ep_size: + raise ValueError("num_experts must be divisible by VLLM EP") + experts_per_rank = num_experts // ep_size + coverage: dict[int, set[tuple[int, int]]] = defaultdict(set) + params = [] + for info in infos: + layer, projection, expert = map(int, _ROUTED_EXPERT.fullmatch(info.name).groups()) + if not 0 <= expert < num_experts: + raise ValueError(f"invalid expert id {expert} in {info.name}") + ep_rank = expert // experts_per_rank + coverage[layer].add((expert, projection)) + params.append( + _ExpertParam( + info=info, + layer=layer, + expert=expert, + target_ranks=tuple(target_ranks[ep_rank]), + ) + ) + + expected = {(expert, projection) for expert in range(num_experts) for projection in (1, 2)} + if not coverage or any(found != expected for found in coverage.values()): + raise ValueError("routed-expert metadata is incomplete") + return sorted(params, key=lambda param: (param.layer, param.info.name)) + + +def _set_expert_source_ranks( + infos: Sequence[ParamInfo], + local_names_by_rank: Sequence[Sequence[str]], +) -> list[ParamInfo]: + owners = {} + for rank, names in enumerate(local_names_by_rank): + for name in names: + owners.setdefault(name, rank) + missing = [info.name for info in infos if info.name not in owners] + if missing: + raise ValueError(f"no physical owner for {missing[0]}") + return [replace(info, src_rank=owners[info.name]) for info in infos] + + +def _resolve_expert_source_ranks( + infos: Sequence[ParamInfo], + get_local_weight_names: Callable[[], Iterable[str]], +) -> list[ParamInfo]: + local_expert_names = tuple(name for name in get_local_weight_names() if _ROUTED_EXPERT.fullmatch(name)) + local_names_by_rank = [None] * dist.get_world_size() + dist.all_gather_object(local_names_by_rank, local_expert_names, group=get_gloo_group()) + return _set_expert_source_ranks(infos, local_names_by_rank) + + +def _build_expert_transfer_plan( + params: Sequence[_ExpertParam], + buffer_size: int, +) -> list[_ExpertTransferGroup]: + """Build expert transfer groups with pre-packed, rank-bounded transfer batches.""" + if buffer_size <= 0: + raise ValueError("update_weight_buffer_size must be positive") + + params_by_transfer: dict[tuple[int, int, tuple[int, ...], int], list[_ExpertParam]] = defaultdict(list) + for param in params: + params_by_transfer[(param.layer, param.expert, param.target_ranks, param.info.src_rank)].append(param) + + by_layer: dict[int, list[_ExpertTransfer]] = defaultdict(list) + for (layer, _expert, target_ranks, source_rank), transfer_params in params_by_transfer.items(): + transfer = _ExpertTransfer( + source_rank=source_rank, + target_ranks=target_ranks, + params=tuple(sorted(transfer_params, key=lambda param: (param.expert, param.info.name))), + ) + by_layer[layer].append(transfer) + + transfer_plan = [] + for layer in sorted(by_layer): + transfer_group = tuple( + sorted(by_layer[layer], key=lambda transfer: (transfer.target_ranks, transfer.source_rank)) + ) + transfer_plan.append(tuple(_pack_expert_transfer_batches(transfer_group, buffer_size))) + return transfer_plan + + +def _expert_transfer_size(transfer: _ExpertTransfer) -> int: + return sum(param.info.size for param in transfer.params) + + +def _pack_expert_transfer_batches( + transfers: Sequence[_ExpertTransfer], + buffer_size: int, +) -> list[_ExpertTransferBatch]: + """First-fit transfers while capping per-rank staging bytes.""" + sized_transfers = sorted( + ((_expert_transfer_size(transfer), transfer) for transfer in transfers), + key=lambda item: (-item[0], item[1].target_ranks, item[1].source_rank), + ) + if buffer_size < sized_transfers[0][0]: + raise ValueError("one source-to-target expert transfer bundle exceeds update_weight_buffer_size") + + batches: list[list[_ExpertTransfer]] = [] + batch_costs: list[dict[int, int]] = [] + for size, transfer in sized_transfers: + participants = set(transfer.target_ranks) | {transfer.source_rank} + candidates = [ + index + for index, costs in enumerate(batch_costs) + if all(costs.get(rank, 0) + size <= buffer_size for rank in participants) + ] + if candidates: + batch_index = min(candidates, key=lambda index: (sum(batch_costs[index].values()), index)) + else: + batch_index = len(batches) + batches.append([]) + batch_costs.append({}) + + batches[batch_index].append(transfer) + for rank in participants: + batch_costs[batch_index][rank] = batch_costs[batch_index].get(rank, 0) + size + + return [tuple(batch) for batch in batches] + + +def _log_disabled_expert_routing(reason: str) -> None: + if dist.get_rank() == 0: + logger.info("Disable rank-local expert update: %s", reason) + + +def configure_expert_routing( + *, + args: Namespace, + full_param_info_buckets: Sequence[Sequence[ParamInfo]] | None, + get_local_weight_names: Callable[[], Iterable[str]], + engine_gpu_counts: Sequence[int], + engine_gpu_offsets: Sequence[int], + engine_parallel_configs: Sequence[Mapping[str, Any]] | None, + use_distribute: bool, +) -> tuple[list[list[ParamInfo]] | None, list[_ExpertTransferGroup]]: + if full_param_info_buckets is None: + return None, [] + + if use_distribute: + _log_disabled_expert_routing("distributed VLLM engines are present") + return None, [] + if not engine_gpu_counts: + _log_disabled_expert_routing("no colocated VLLM engines") + return None, [] + + try: + vllm_moe_topology = _get_homogeneous_vllm_moe_topology( + args, + engine_gpu_counts, + engine_parallel_configs, + ) + except (AttributeError, TypeError, ValueError) as exc: + _log_disabled_expert_routing(str(exc)) + return None, [] + + if not _can_route_experts( + args, + vllm_moe_topology, + engine_gpu_counts=engine_gpu_counts, + ): + _log_disabled_expert_routing("VLLM/Megatron expert topology is not eligible") + return None, [] + dense_infos = [] + expert_infos = [] + for bucket in full_param_info_buckets: + for info in bucket: + (expert_infos if _ROUTED_EXPERT.fullmatch(info.name) else dense_infos).append(info) + if not expert_infos: + return None, [] + + try: + from megatron.core import mpu + + from .hf_weight_iterator_direct import pack_param_info_buckets + + expert_infos = _resolve_expert_source_ranks(expert_infos, get_local_weight_names) + target_ranks = _get_expert_target_ranks( + engine_gpu_counts, + engine_gpu_offsets, + ep_size=vllm_moe_topology.ep_size, + world_size=dist.get_world_size(), + ) + expert_params = _build_expert_params( + expert_infos, + target_ranks, + num_experts=args.num_experts, + ) + buffer_size = args.update_weight_buffer_size + expert_transfer_plan = _build_expert_transfer_plan(expert_params, buffer_size) + expert_transfer_batches = sum(len(group) for group in expert_transfer_plan) + dense_buckets = pack_param_info_buckets(dense_infos, buffer_size) + except (AttributeError, TypeError, ValueError) as exc: + _log_disabled_expert_routing(str(exc)) + return None, [] + + if dist.get_rank() == 0: + logger.info( + "Enabled rank-local expert update: Megatron PP=%d EP=%d, VLLM EP=%d, " + "%d -> %d transfer groups (%d dense + %d expert, %d expert transfer batches)", + mpu.get_pipeline_model_parallel_world_size(), + mpu.get_expert_model_parallel_world_size(), + vllm_moe_topology.ep_size, + len(full_param_info_buckets), + len(dense_buckets) + len(expert_transfer_plan), + len(dense_buckets), + len(expert_transfer_plan), + expert_transfer_batches, + ) + return dense_buckets, expert_transfer_plan diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py deleted file mode 100644 index 369f8c2d4..000000000 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py +++ /dev/null @@ -1,29 +0,0 @@ -from abc import ABC, abstractmethod - - -class HfWeightIteratorBase(ABC): - @staticmethod - def create(args, model, **kwargs): - from .hf_weight_iterator_bridge import HfWeightIteratorBridge - from .hf_weight_iterator_direct import HfWeightIteratorDirect - - c = { - "raw": HfWeightIteratorDirect, - "bridge": HfWeightIteratorBridge, - }[args.megatron_to_hf_mode] - - return c(args, model, **kwargs) - - def __init__(self, args, model, model_name, quantization_config): - self.args = args - self.model = model - self.model_name = model_name - self.quantization_config = quantization_config - - @abstractmethod - def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): - """ - Mental model of the API: - megatron_model.to_hf_magically().named_parameters() - """ - raise NotImplementedError diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py deleted file mode 100644 index d9aa7338f..000000000 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py +++ /dev/null @@ -1,112 +0,0 @@ -import dataclasses - - -from vime.utils import megatron_bridge_utils -from vime.utils.misc import chunk_named_params_by_size - -from ..megatron_to_hf import postprocess_hf_param -from ..megatron_to_hf.processors import quantize_params -from ..misc_utils import strip_param_name_prefix -from .hf_weight_iterator_base import HfWeightIteratorBase - - -def _patch_bridge_expert_cache_to_cpu(): - """Monkey-patch GPTOSSBridge class to cache expert weights on CPU. - - This avoids GPU OOM when torch.cat merges all experts, especially in - colocated mode where vLLM and Megatron share the same GPU. - """ - try: - from megatron.bridge.models.gpt_oss.gpt_oss_bridge import GPTOSSBridge - except ImportError: - return - - if getattr(GPTOSSBridge, "_cpu_cache_patched", False): - return - - _orig = GPTOSSBridge.maybe_modify_converted_hf_weight - - def _patched(self, task, converted_weights_dict): - cpu_dict = {k: v.cpu() for k, v in converted_weights_dict.items()} - result = _orig(self, task, cpu_dict) - # Move merged result back to GPU for CUDA IPC serialization - return {k: v.cuda() for k, v in result.items()} if result else result - - GPTOSSBridge.maybe_modify_converted_hf_weight = _patched - GPTOSSBridge._cpu_cache_patched = True - - -class HfWeightIteratorBridge(HfWeightIteratorBase): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - from megatron.bridge import AutoBridge - - import vime_plugins.megatron_bridge # noqa: F401 - - self._bridge = megatron_bridge_utils.patch_auto_bridge_hf_config( - AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True) - ) - _patch_bridge_expert_cache_to_cpu() - - def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): - # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name) - renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()} - with megatron_bridge_utils.patch_megatron_model(self.model): - conversion_tasks = self._bridge.get_conversion_tasks(self.model) - conversion_tasks = _process_conversion_tasks(conversion_tasks, renamed_megatron_local_weights) - - named_weights = self._bridge.export_hf_weights(self.model, cpu=False, conversion_tasks=conversion_tasks) - - def _streaming_quantized(): - for hf_param_name, weight, megatron_param_name in named_weights: - processed_weight = postprocess_hf_param( - args=self.args, - megatron_param_name=megatron_param_name, - hf_param_name=hf_param_name, - param=weight, - ) - converted_named_params = [(hf_param_name, processed_weight)] - quantized_batch = quantize_params( - args=self.args, - megatron_name=megatron_param_name, - converted_named_params=converted_named_params, - quantization_config=self.quantization_config, - ) - yield from quantized_batch - - yield from chunk_named_params_by_size( - _streaming_quantized(), chunk_size=self.args.update_weight_buffer_size - ) - - -def _process_conversion_tasks(vanilla_conversion_tasks, new_weight_dict): - def _handle_one(task): - if task is None: - return None - if task.param_weight is None: - return task - - weight_dict_key = f"vp_stages.{task.vp_stage}.{task.param_name}" - assert ( - weight_dict_key in new_weight_dict - ), f"{weight_dict_key=} not in new_weight_dict ({task.vp_stage=}, {task.param_name=}, {list(new_weight_dict)=})" - - new_param_weight = new_weight_dict[weight_dict_key] - new_param_weight = new_param_weight.cuda() - return dataclasses.replace(task, param_weight=new_param_weight) - - return _MapWithLen(_handle_one, vanilla_conversion_tasks) - - -class _MapWithLen: - def __init__(self, fn, xs): - self.fn = fn - self.xs = xs - - def __len__(self): - return len(self.xs) - - def __iter__(self): - for x in self.xs: - yield self.fn(x) diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index d345adde8..160868e68 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -1,41 +1,79 @@ import dataclasses from argparse import Namespace -from collections.abc import Sequence +from collections.abc import Callable, Sequence import torch import torch.distributed as dist from megatron.core import mpu from tqdm import tqdm +from vime.utils import accelerator from vime.utils.distributed_utils import get_gloo_group from vime.utils.types import ParamInfo from ..megatron_to_hf import convert_to_hf from .common import all_gather_params_async, named_params_and_buffers -from .hf_weight_iterator_base import HfWeightIteratorBase -class HfWeightIteratorDirect(HfWeightIteratorBase): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) +class HfWeightIteratorDirect: + def __init__(self, args, model, model_name, quantization_config, transform_ue8m0=False): + self.args = args + self.model = model + self.model_name = model_name + self.quantization_config = quantization_config + self.transform_ue8m0 = transform_ue8m0 self.megatron_local_param_info_buckets = _get_megatron_local_param_info_buckets(self.args, self.model) + self.ep_broadcast_src_rank_map = _get_ep_broadcast_src_rank_map() - def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): + def get_hf_weight_chunks( + self, + megatron_local_weights, + progress_desc: str = "Update weights", + should_convert_chunk: Callable[[int], bool] | None = None, + param_info_buckets: Sequence[Sequence[ParamInfo]] | None = None, + ): rank = dist.get_rank() + param_info_buckets = ( + self.megatron_local_param_info_buckets if param_info_buckets is None else param_info_buckets + ) - for megatron_local_param_infos in tqdm( - self.megatron_local_param_info_buckets, disable=rank != 0, desc=progress_desc + for chunk_idx, megatron_local_param_infos in enumerate( + tqdm(param_info_buckets, disable=rank != 0, desc=progress_desc) ): - megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights) - hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos) - yield hf_named_tensors - del megatron_full_params + megatron_full_params = _get_megatron_full_params( + megatron_local_param_infos, + megatron_local_weights, + self.args.update_weight_buffer_size, + self.ep_broadcast_src_rank_map, + ) + if should_convert_chunk is None or should_convert_chunk(chunk_idx): + hf_named_tensors = self._convert_to_hf_named_tensors( + megatron_full_params, + megatron_local_param_infos, + ) + else: + hf_named_tensors = [] + try: + yield hf_named_tensors + finally: + del hf_named_tensors, megatron_full_params - def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tensor], param_infos: list[ParamInfo]): + def _convert_to_hf_named_tensors( + self, + megatron_full_params: Sequence[torch.Tensor], + param_infos: Sequence[ParamInfo], + ): hf_named_tensors = [] for info, param in zip(param_infos, megatron_full_params, strict=False): hf_named_tensors.extend( - convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config) + convert_to_hf( + self.args, + self.model_name, + info.name, + param, + self.quantization_config, + transform_ue8m0=self.transform_ue8m0, + ) ) return hf_named_tensors @@ -43,24 +81,24 @@ def _convert_to_hf_named_tensors(self, megatron_full_params: Sequence[torch.Tens def _get_megatron_full_params( megatron_local_param_infos: Sequence[ParamInfo], megatron_local_weights, + broadcast_buffer_size: int, + ep_broadcast_src_rank_map: dict[int, int], ) -> Sequence[torch.Tensor]: pp_size = mpu.get_pipeline_model_parallel_world_size() ep_size = mpu.get_expert_model_parallel_world_size() - rank = dist.get_rank() # init params: params = [] for info in megatron_local_param_infos: if dist.get_rank() == info.src_rank: params.append( torch.nn.Parameter( - megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True), + megatron_local_weights[info.name].to(device=accelerator.current_device(), non_blocking=True), requires_grad=False, ) ) else: - params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device())) - torch.cuda.synchronize() - + params.append(torch.empty(info.shape, dtype=info.dtype, device=accelerator.current_device())) + accelerator.synchronize() # broadcast params across pp ranks if pp_size > 1: handles = [] @@ -76,21 +114,7 @@ def _get_megatron_full_params( # broadcast params across ep ranks if ep_size > 1: - handles = [] - for info, param in zip(megatron_local_param_infos, params, strict=False): - if ".experts." in info.name: - src_rank = ( - info.src_rank - if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group()) - else rank - ) - handles.append( - torch.distributed.broadcast( - param, src=src_rank, group=mpu.get_expert_model_parallel_group(), async_op=True - ) - ) - for handle in handles: - handle.wait() + _broadcast_expert_params(megatron_local_param_infos, params, broadcast_buffer_size, ep_broadcast_src_rank_map) # Set tp attrs for all params for info, param in zip(megatron_local_param_infos, params, strict=False): @@ -103,11 +127,53 @@ def _get_megatron_full_params( return gathered_params +def _broadcast_expert_params( + param_infos: Sequence[ParamInfo], + params: Sequence[torch.Tensor], + buffer_size: int, + src_rank_map: dict[int, int], +) -> None: + ep_group = mpu.get_expert_model_parallel_group() + params_by_src: dict[int, list[torch.Tensor]] = {} + for info, param in zip(param_infos, params, strict=False): + if ".experts." not in info.name: + continue + params_by_src.setdefault(src_rank_map[info.src_rank], []).append(param) + + for src_rank, expert_params in params_by_src.items(): + dist._broadcast_coalesced(ep_group, expert_params, buffer_size, src=src_rank) + + +def _get_ep_broadcast_src_rank_map() -> dict[int, int]: + ep_group = mpu.get_expert_model_parallel_group() + ep_size = mpu.get_expert_model_parallel_world_size() + if ep_size == 1: + return {dist.get_rank(): 0} + + pp_group_ranks = dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()) + pp_groups: list[list[int] | None] = [None] * ep_size + dist.all_gather_object(pp_groups, pp_group_ranks, group=ep_group) + + src_rank_map = {} + for ep_rank, pp_group in enumerate(pp_groups): + assert pp_group is not None + for global_rank in pp_group: + src_rank_map[global_rank] = ep_rank + return src_rank_map + + def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torch.nn.Module]) -> list[list[ParamInfo]]: """ Partition params into buckets ≤ update_weight_buffer_size (with TP replication). """ param_infos = _get_megatron_local_param_infos(args, model) + return pack_param_info_buckets(param_infos, args.update_weight_buffer_size) + + +def pack_param_info_buckets( + param_infos: Sequence[ParamInfo], + update_weight_buffer_size: int, +) -> list[list[ParamInfo]]: param_info_buckets = [[]] # Start with one empty bucket buffer_size = 0 # Track current bucket size in bytes @@ -122,7 +188,7 @@ def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torc param_size = info.size * tp_size # If adding this param exceeds limit AND current bucket has params: start new bucket - if buffer_size + param_size > args.update_weight_buffer_size and len(param_info_buckets[-1]) > 0: + if buffer_size + param_size > update_weight_buffer_size and len(param_info_buckets[-1]) > 0: param_info_buckets.append([]) buffer_size = 0 @@ -144,6 +210,9 @@ def _get_megatron_local_param_infos(args: Namespace, model: Sequence[torch.nn.Mo param_infos = {} rank = dist.get_rank() for name, param in named_params_and_buffers(args, model): + # DSpark draft params are sent through the draft weight source. + if ".draft_model." in name: + continue param_infos[name] = ParamInfo( name=name, dtype=param.dtype, diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py new file mode 100644 index 000000000..1f0cc8559 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import shutil +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path + +import torch +import torch.distributed as dist +from ray.actor import ActorHandle + +from vime.utils.distributed_utils import get_gloo_group + +from ..hf_checkpoint_saver import save_hf_model_to_path + + +class UpdateWeightFromDisk: + """Full-weight sync through a shared filesystem and VLLM disk reload.""" + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.model = model + self.weights_getter = weights_getter + self.model_name = model_name + self.quantization_config = quantization_config + self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} + self.rollout_engines: Sequence[ActorHandle] = [] + self.rollout_engine_lock: ActorHandle | None = None + # Post-write hook: object-store-backed shared filesystems lack cross-host + # read-after-write consistency, so written files need an explicit step + # (e.g. uploading them to the backing object store) before the engines can see them. + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from vime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + engine_parallel_configs: Sequence[Mapping[str, object]] | None = None, + ) -> None: + self.rollout_engines = rollout_engines + self.rollout_engine_lock = rollout_engine_lock + + def disconnect_rollout_engines(self) -> None: + return + + def pop_metrics(self) -> dict[str, float]: + out, self.update_weight_metrics = self.update_weight_metrics, {} + return out + + @torch.no_grad() + def update_weights(self) -> None: + self.weight_version += 1 + version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{self.weight_version:06d}" + + if dist.get_rank() == 0: + shutil.rmtree(version_dir, ignore_errors=True) + dist.barrier(group=get_gloo_group()) + + # every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface + # one rank's mkdir to another until commit + version_dir.mkdir(parents=True, exist_ok=True) + save_hf_model_to_path( + self.args, + version_dir, + self.model, + model_name=self.model_name, + quantization_config=self.quantization_config, + progress_desc="Save HF weights for update from disk", + ) + dist.barrier(group=get_gloo_group()) + + # every rank runs the hook (it gates itself): each container must publish + # its own writes + if self._post_write_hook is not None: + self._post_write_hook(self.args, str(version_dir), list(self.rollout_engines)) + dist.barrier(group=get_gloo_group()) + + # VLLM reload is orchestrated by RayTrainGroup after the checkpoint + # is fully written, so training-side lifecycle can decide whether + # Megatron actors are still alive. diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py new file mode 100644 index 000000000..d3725dffb --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import json +import logging +import os +import queue +import shutil +from argparse import Namespace +from collections import deque +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import ray +import safetensors.numpy +import torch +import torch.distributed as dist +import zstandard +from ray.actor import ActorHandle + +from vime.utils import accelerator +from vime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode +from vime.utils.distributed_utils import get_gloo_group + +from .update_weight_from_distributed import UpdateWeightFromDistributed + +logger = logging.getLogger(__name__) + + +class UpdateWeightFromDiskDelta(UpdateWeightFromDistributed): + """ + Delta weight sync over a shared filesystem. PP-src ranks diff each gathered HF tensor against + a CPU snapshot of the previous sync and publish the changes as a canonical HF checkpoint dir; + each engine's /pull_weights fans the apply out to every host it spans, then the engine reloads + the patched local checkpoint via the ordinary update_weights_from_disk path. vime only ever + talks to one endpoint per engine, so multi-node serving and external engines need nothing extra. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + super().__init__(args, model, weights_getter, model_name=model_name, quantization_config=quantization_config) + self.delta_dir = args.update_weight_disk_dir + os.makedirs(self.delta_dir, exist_ok=True) + self.delta_encoding = args.update_weight_delta_encoding + self.checksum_algorithm = args.update_weight_delta_checksum + self._snapshot: dict[str, np.ndarray] = {} + self._baseline_captured = False + # Post-write hook: object-store-backed shared filesystems lack cross-host + # read-after-write consistency, so written files need an explicit step + # (e.g. uploading them to the backing object store) before the engines can see them. + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from vime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + engine_parallel_configs: Sequence[Mapping[str, object]] | None = None, + ) -> None: + # The rollout_engine_lock the NCCL path uses isn't needed — the engine-side apply is + # serialized by a per-host flock. + self.rollout_engines = rollout_engines + self._is_pp_src_rank = dist.get_rank() == 0 + + def disconnect_rollout_engines(self) -> None: + pass # no NCCL groups to tear down + + @torch.no_grad() + def update_weights(self) -> None: + # The first call only captures the baseline snapshot the next sync diffs against. + if not self._baseline_captured: + self._capture_baseline() + self._baseline_captured = True + return + + self.weight_version += 1 + self._publish() + self._reload_engines() + self._record_metrics() + + def _capture_baseline(self) -> None: + """Capture the baseline snapshot the first delta diffs against (no publish), and clear any + stale stream from a prior run. Seeds from hf_checkpoint — what each host materializes its + base from — so the invariant ``snapshot == engine base`` holds even where the megatron->HF + round-trip trims vocab-padding rows (embed/lm_head). A tensor absent there (rare) falls back + to the gathered value. pull_weights(0) makes each host materialize its local base now, + overlapped with the snapshot gather, so the first real sync only pays the delta apply.""" + # a prior run's versions would apply against the wrong base; start the dir clean + pulls = [] + if dist.get_rank() == 0: + shutil.rmtree(self.delta_dir, ignore_errors=True) + os.makedirs(self.delta_dir, exist_ok=True) + if self._post_write_hook is not None: + self._post_write_hook(self.args, self.delta_dir, list(self.rollout_engines)) + pulls = [engine.pull_weights.remote(target_version=0) for engine in self.rollout_engines] + dist.barrier(group=get_gloo_group()) + + read_hf = make_tensor_reader(self.args.hf_checkpoint) # index the HF headers once + for name, tensor in self._iter_hf_tensors(): + try: + self._snapshot[name] = read_hf(name) + except KeyError: + self._snapshot[name] = tensor.detach().cpu().contiguous().view(torch.uint8).numpy().reshape(-1) + logger.warning("seed: %s absent from hf_checkpoint; seeding from current weights", name) + if dist.get_rank() == 0: + ray.get(pulls) + logger.info( + "[disk delta] captured baseline snapshot of %d tensors from %s", + len(self._snapshot), + self.args.hf_checkpoint, + ) + + def _publish(self) -> None: + """Encode this version's changed tensors (PP-src ranks), then write it as a canonical HF dir.""" + self._encode_delta() + dist.barrier(group=get_gloo_group()) + self._write_delta_files() + + def _write_delta_files(self) -> None: + """Write this rank's changed tensors as one canonical model-NNNNN.safetensors, and on rank + 0 the HF index. The sequential file numbers and the index are coordinated over gloo (small + object gathers), not the filesystem — a non-POSIX shared filesystem may not surface one rank's writes to + another until commit.""" + group = get_gloo_group() + world, rank = dist.get_world_size(), dist.get_rank() + + # number the files sequentially across only the ranks that have one (no gaps) + counts: list = [None] * world + dist.all_gather_object(counts, int(bool(self._delta)), group=group) + offset, total = sum(counts[:rank]), sum(counts) + + fname = None + self.wire_bytes = 0 + if self._delta: + fname = f"model-{offset:05d}-of-{total:05d}.safetensors" + blob = safetensors.numpy.save(self._delta, metadata=self._checksums) + self.wire_bytes = len(blob) + _atomic_write(os.path.join(self._version_dir, fname), blob) + + maps: list = [None] * world + dist.all_gather_object(maps, {name: fname for name in self._delta}, group=group) + if rank == 0: + index = { + "metadata": { + "version": f"{self.weight_version:06d}", + "base_version": f"{self.weight_version - 1:06d}", + "delta_encoding": self.delta_encoding, + "compression_format": "zstd", + "checksum_format": self.checksum_algorithm, + }, + "weight_map": {name: f for m in maps for name, f in m.items()}, + } + _atomic_write(os.path.join(self._version_dir, "model.safetensors.index.json"), json.dumps(index).encode()) + dist.barrier(group=group) + + def _reload_engines(self) -> None: + """Commit the published files, have each engine pull the delta onto every host it spans + (checksum-verified), then reload the engines.""" + if self._post_write_hook is not None: + self._post_write_hook(self.args, self._version_dir, list(self.rollout_engines)) + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + ray.get([engine.pull_weights.remote(self.weight_version) for engine in self.rollout_engines]) + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + ray.get( + [ + engine.update_weights_from_disk.remote( + model_path=self.args.update_weight_local_checkpoint_dir, + weight_version=str(self.weight_version), + ) + for engine in self.rollout_engines + ] + ) + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + def _iter_hf_tensors(self): + """Yield gathered HF tensors on the publishing rank.""" + for name, tensor in self._source: + if self._is_pp_src_rank: + yield name, tensor + + def _encode_delta(self) -> None: + """Diff each gathered HF tensor against the snapshot, keeping the changed ones (compressed) + in self._delta with their checksums. The GPU->CPU gather is pipelined into a compute pool: + the main loop copies one tensor to a pinned buffer and submits it; pool workers diff and + compress in parallel (each is a few big GIL-releasing numpy/zstd calls).""" + self._version_dir = os.path.join(self.delta_dir, f"weight_v{self.weight_version:06d}") + if self._is_pp_src_rank: + os.makedirs(self._version_dir, exist_ok=True) + snapshot = self._snapshot + self._delta: dict[str, np.ndarray] = {} # changed tensor name -> compressed diff + self._checksums: dict[str, str] = {} # changed tensor name -> new-state checksum + self.changed_bytes = self.total_bytes = 0 + + # Pinned host-buffer pool: a pinned non_blocking GPU->CPU copy is far faster than .cpu(). + max_bytes = max((int(v.nbytes) for v in snapshot.values()), default=0) + free_q: queue.Queue = queue.Queue() + use_pinned = True + try: + for _ in range(max(4, min(2 * NUM_WORKERS, (32 << 30) // max(max_bytes, 1)))): + free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) + except RuntimeError as e: # low memlock limit + logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e) + use_pinned = False + + def diff_and_compress(name, buf, nbytes, pinned): + if pinned: # copy out and free the pinned buffer before the heavy diff/compress + new = np.empty(nbytes, dtype=np.uint8) + np.copyto(new, buf.numpy()[:nbytes]) + free_q.put(buf) + else: + new = buf + old = snapshot[name] + if self.delta_encoding == "xor": + diff = new ^ old + changed = int(np.count_nonzero(diff)) + elif self.delta_encoding == "overwrite": + mask = new != old + changed = int(np.count_nonzero(mask)) + diff = overwrite_encode(new, mask) + else: + raise ValueError(f"unknown delta encoding {self.delta_encoding!r}") + if not changed: + return name, new, None, None, 0 + compressed = np.frombuffer(zstandard.ZstdCompressor(level=1).compress(diff), dtype=np.uint8) + return name, new, compressed, checksum(self.checksum_algorithm, new), changed + + def collect(fut): + name, new, compressed, digest, changed = fut.result() + snapshot[name] = new # becomes the next sync's base + if changed: + self.changed_bytes += changed + self._delta[name] = compressed + self._checksums[name] = digest + + pool = ThreadPoolExecutor(max_workers=NUM_WORKERS) + inflight: deque = deque() + try: + for name, tensor in self._iter_hf_tensors(): + flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) + nbytes = int(flat.numel()) + if use_pinned and nbytes <= max_bytes: + buf = free_q.get() # blocks when all buffers are in flight -> backpressures the gather + buf[:nbytes].copy_(flat, non_blocking=True) + accelerator.current_stream().synchronize() + payload, pinned = buf, True + else: + payload, pinned = flat.cpu().numpy(), False + self.total_bytes += nbytes + inflight.append(pool.submit(diff_and_compress, name, payload, nbytes, pinned)) + if len(inflight) >= 2 * NUM_WORKERS: + collect(inflight.popleft()) + while inflight: + collect(inflight.popleft()) + finally: + pool.shutdown() + + def _record_metrics(self) -> None: + """All-reduce the byte counts and record changed-fraction / wire size; the actor drains + update_weight_metrics onto the step log.""" + counts = torch.tensor( + [self.changed_bytes, self.total_bytes, self.wire_bytes], + dtype=torch.int64, + device=accelerator.current_device(), + ) + dist.all_reduce(counts) + changed, total, wire = counts.tolist() + m = self.update_weight_metrics + m["perf/update_weights_density"] = changed / max(total, 1) + m["perf/update_weights_wire_bytes"] = wire + if dist.get_rank() == 0: + logger.info( + "[disk delta v=%s] density=%.2f%% wire=%.2f GB", + self.weight_version, + 100.0 * changed / max(total, 1), + wire / 1e9, + ) + + +def _atomic_write(path: str, data: bytes) -> None: + tmp = path + ".tmp" + with open(tmp, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index dd7ff1bb1..6919d8d3f 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,52 +1,21 @@ -from __future__ import annotations - -import logging -import os -import socket -import time from argparse import Namespace -from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any +from collections.abc import Callable, Mapping, Sequence +from functools import partial import ray import torch import torch.distributed as dist -from megatron.core import mpu -from ray import ObjectRef from ray.actor import ActorHandle -from tqdm import tqdm -from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine -from vllm_ascend.distributed.weight_transfer.hccl_engine import HCCLTrainerSendWeightsArgs, HCCLWeightTransferEngine -from vime.utils.common import is_npu from vime.utils.distributed_utils import get_gloo_group -from ..megatron_to_hf import convert_to_hf -from .common import all_gather_param, named_params_and_buffers -from .hf_weight_iterator_base import HfWeightIteratorBase - -logger = logging.getLogger(__name__) - - -def _begin_vllm_weight_update_session(rollout_engines: Sequence[ActorHandle]) -> None: - if dist.get_rank() == 0: - logger.info("vLLM weight update: start_weight_update") - ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in rollout_engines]) - dist.barrier(group=get_gloo_group()) - - -def _end_vllm_weight_update_session(rollout_engines: Sequence[ActorHandle]) -> None: - if dist.get_rank() == 0: - logger.info("vLLM weight update: finish_weight_update") - ray.get([engine.finish_weight_update.remote() for engine in rollout_engines]) - dist.barrier(group=get_gloo_group()) +from ..dspark.export import export_dspark_model_weights +from .common import HfWeightSource, VimeRayWeightSyncClient, create_nccl_trainer +from .hf_weight_iterator_direct import HfWeightIteratorDirect class UpdateWeightFromDistributed: - """ - Update distributed engines via NCCL. Each PP rank: group "vime-pp_{pp_rank}", - only DP=TP=0 broadcasts. Non-expert (TP) and expert (EP) params separate. - """ + """Update distributed vLLM engines through its stateful NCCL trainer API.""" def __init__( self, @@ -57,26 +26,27 @@ def __init__( model_name: str, quantization_config: dict[str, int | str | list[str]] | None, ) -> None: - """ - Initialize. Groups created in connect_rollout_engines. - """ self.args = args - self.model = model - self.weights_getter = weights_getter - self.model_name = model_name self.quantization_config = quantization_config self.weight_version = 0 - self._model_update_groups = None - self._hf_weight_iterator = ( - HfWeightIteratorBase.create( - args=args, - model=model, - model_name=model_name, - quantization_config=quantization_config, + self.update_weight_metrics: dict[str, float] = {} + iterator = HfWeightIteratorDirect( + args=args, + model=model, + model_name=model_name, + quantization_config=quantization_config, + ) + draft_weights_getter = ( + partial( + export_dspark_model_weights, + model, + use_policy_embedding=not self.args.dspark_pretrained_model, ) - if args.megatron_to_hf_mode == "bridge" + if self.args.dspark_enabled else None ) + self._source = HfWeightSource(iterator, weights_getter, draft_weights_getter) + self._trainer = None def connect_rollout_engines( self, @@ -84,56 +54,42 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, + engine_parallel_configs: Sequence[Mapping[str, object]] | None = None, ) -> None: - """ - Create NCCL "vime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent broadcasts. - """ - self.rollout_engines = rollout_engines - self.rollout_engine_lock = rollout_engine_lock - self._engine_gpu_counts = engine_gpu_counts - - # For TP: - # 1. AllGather parameters to rank 0 - # 2. Broadcast parameters from rank 0 to all vLLM engines - self._is_pp_src_rank = ( - mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 + del rollout_engine_lock, engine_gpu_offsets, engine_parallel_configs + self.disconnect_rollout_engines() + self.rollout_engines = list(rollout_engines) + engine_gpu_counts = list( + engine_gpu_counts or [self.args.rollout_num_gpus_per_engine] * len(self.rollout_engines) + ) + client = VimeRayWeightSyncClient( + self.rollout_engines, + lambda: self.weight_version, + engine_gpu_counts, + ) + self._trainer = create_nccl_trainer( + client, + self._source, + engine_gpu_counts, ) - pp_rank = mpu.get_pipeline_model_parallel_rank() - if self._is_pp_src_rank: - self._group_name = f"vime-pp_{pp_rank}" - - if self._is_pp_src_rank: - if self._model_update_groups is not None: - disconnect_rollout_engines_from_distributed( - self.args, self._group_name, self._model_update_groups, self.rollout_engines - ) - self._model_update_groups = connect_rollout_engines_from_distributed( - self.args, - self._group_name, - rollout_engines, - engine_gpu_counts=engine_gpu_counts, - ) def disconnect_rollout_engines(self) -> None: - if not getattr(self, "_is_pp_src_rank", False) or self._model_update_groups is None: - return - disconnect_rollout_engines_from_distributed( - self.args, self._group_name, self._model_update_groups, self.rollout_engines - ) - self._model_update_groups = None + if self._trainer is not None: + self._trainer.shutdown() + self._trainer = None + + def pop_metrics(self) -> dict[str, float]: + metrics, self.update_weight_metrics = self.update_weight_metrics, {} + return metrics @torch.no_grad() def update_weights(self) -> None: - """ - Pause → flush → _send_weights → continue. Progress on PP source. - """ + assert self._trainer is not None self.weight_version += 1 if dist.get_rank() == 0: ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) - - # int4/fp4 pre_process if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: post_process_weights( restore_weights_before_load=True, @@ -142,18 +98,20 @@ def update_weights(self) -> None: ) dist.barrier(group=get_gloo_group()) - pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None - _begin_vllm_weight_update_session(self.rollout_engines) - try: - self._send_weights(pbar) - if self._is_pp_src_rank: - torch.cuda.synchronize() - finally: - _end_vllm_weight_update_session(self.rollout_engines) + client = self._trainer.client + client.draft = False + self._trainer.send_weights() + update_draft = self.args.dspark_enabled or ( + self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp" + ) + if update_draft: + self._source.draft = self.args.dspark_enabled + client.draft = True + self._trainer.send_weights() + self._source.draft = False + client.draft = False - dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: - # int4/fp4 post_process if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: post_process_weights( restore_weights_before_load=False, @@ -163,345 +121,12 @@ def update_weights(self) -> None: ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) - def _send_weights(self, pbar: tqdm | None) -> None: - """ - Non-expert (TP) pass → barrier → expert (EP) pass → barrier. Each iterator - yields broadcast-ready chunks (bucketing happens internally). - """ - use_vllm_packed = self._use_vllm_packed() - if self._hf_weight_iterator is not None: - self._sync_bridge_weights_to_rollout_engines(pbar, use_vllm_packed=use_vllm_packed) - return - - if use_vllm_packed and self._is_pp_src_rank: - logger.info("Using vLLM packed weight sync (bucketed; metadata + trainer_send_weights per bucket)") - - for hf_chunk in self._iter_non_expert_chunks(): - self._update_bucket_weights_from_distributed(hf_chunk, pbar=pbar, packed=use_vllm_packed) - dist.barrier(group=get_gloo_group()) - - if not use_vllm_packed: - for hf_chunk in self._iter_expert_chunks(): - self._update_bucket_weights_from_distributed(hf_chunk, pbar=pbar, packed=False) - dist.barrier(group=get_gloo_group()) - - def _sync_bridge_weights_to_rollout_engines(self, pbar: tqdm | None, *, use_vllm_packed: bool) -> None: - """ - Export HF weights through Megatron-Bridge, then send each exported chunk - over the same NCCL non-colocate path used by the raw converter. - """ - if self._is_pp_src_rank: - logger.info("Using Megatron-Bridge HF weight export for non-colocate vLLM weight sync") - - megatron_local_weights = self.weights_getter() - for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): - if self._is_pp_src_rank: - hf_named_tensors = list(hf_named_tensors) - self._update_bucket_weights_from_distributed(hf_named_tensors, pbar=pbar, packed=use_vllm_packed) - - dist.barrier(group=get_gloo_group()) - - def _use_vllm_packed(self) -> bool: - """Use vLLM packed weight transfer (one-shot metadata + trainer_send_weights).""" - if not getattr(self.args, "vllm_weight_sync_packed", True): - return False - if any(".experts." in name for name, _ in named_params_and_buffers(self.args, self.model)): - return False - if self.quantization_config and self.quantization_config.get("quant_method") == "compressed-tensors": - return False - return True - - def _iter_non_expert_chunks(self) -> Iterator[list[tuple[str, torch.Tensor]]]: - """ - Yield broadcast-sized HF chunks of non-expert params: TP all-gather + - HF convert per param, then bucket up to ``--update-weight-buffer-size``. - Empty on non-PP-src ranks (they still join all_gather_param). - """ - buffer_size = 0 - buffer: list[tuple[str, torch.Tensor]] = [] - for name, param in named_params_and_buffers(self.args, self.model): - if ".experts." in name: - continue - param = all_gather_param(name, param) - if not self._is_pp_src_rank: - continue - hf_chunk = convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) - chunk_bytes = sum(t.numel() * t.element_size() for _, t in hf_chunk) - if buffer and buffer_size + chunk_bytes > self.args.update_weight_buffer_size: - yield buffer - buffer = [] - buffer_size = 0 - buffer.extend(hf_chunk) - buffer_size += chunk_bytes - if buffer: - yield buffer - - def _iter_expert_chunks( - self, - params: Iterator[tuple[str, torch.Tensor]] | None = None, - ) -> Iterator[list[tuple[str, torch.Tensor]]]: - """ - Yield one HF chunk per EP-weighted batch of expert params: TP gather + - buffer until threshold, then EP gather + HF convert. - """ - if params is None: - params = ((n, p) for n, p in named_params_and_buffers(self.args, self.model) if ".experts." in n) - - buffer_size = 0 - batch: list[tuple[str, torch.Tensor]] = [] - for name, param in params: - param = all_gather_param(name, param) - param_size = param.numel() * param.element_size() - if ( - buffer_size + param_size - ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size: - hf_chunk = self._ep_gather_and_convert(batch) - if hf_chunk: - yield hf_chunk - batch = [] - buffer_size = 0 - - batch.append((name, param)) - buffer_size += param_size - - if batch: - hf_chunk = self._ep_gather_and_convert(batch) - if hf_chunk: - yield hf_chunk - - def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]]) -> list[tuple[str, torch.Tensor]]: - """ - EP all-gather a buffered batch + HF convert on PP source. Returns HF tensors on - PP source, [] elsewhere. Clears ``named_tensors``. - """ - names = [name for name, _ in named_tensors] - all_names = [None] * mpu.get_expert_model_parallel_world_size() - dist.all_gather_object(all_names, names, group=mpu.get_expert_model_parallel_group()) - - for names in all_names: - assert len(named_tensors) == len(names), f"mismatch names length: {len(named_tensors)} != {len(names)}" - - all_gathered_params = [[] for _ in range(mpu.get_expert_model_parallel_world_size())] - handles = [] - for i, (_name, param) in enumerate(named_tensors): - params = [ - torch.empty_like(param.data, device=torch.cuda.current_device()) - for _ in range(mpu.get_expert_model_parallel_world_size()) - ] - handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True) - handles.append(handle) - for ep_rank, names in enumerate(all_names): - all_gathered_params[ep_rank].append((names[i], params[ep_rank])) - for handle in handles: - handle.wait() - - named_tensors.clear() - if not self._is_pp_src_rank: - return [] - - all_gathered_params = sum(all_gathered_params, []) - converted_hf_tensors = [] - for name, param in all_gathered_params: - converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) - - return converted_hf_tensors - - def _update_bucket_weights_from_distributed( - self, - converted_named_tensors: list[tuple[str, torch.Tensor]], - pbar: tqdm | None = None, - *, - packed: bool = False, - ) -> None: - """ - Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock. - """ - # lock the rollout engines to prevent dead lock on broadcast. - while not ray.get(self.rollout_engine_lock.acquire.remote()): - time.sleep(0.1) - - refs = update_weights_from_distributed( - self._group_name, - self._model_update_groups, - self.weight_version, - self.rollout_engines, - converted_named_tensors, - packed=packed, - ) - - ray.get(refs) - converted_named_tensors.clear() - ray.get(self.rollout_engine_lock.release.remote()) - pbar.update(1) - - -def connect_rollout_engines_from_distributed( - args: Namespace, - group_name: str, - rollout_engines: Sequence[ActorHandle], - engine_gpu_counts: Sequence[int] | None = None, -) -> Any: - """ - Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined. - - ``engine_gpu_counts`` gives the number of GPUs per engine. When engines - have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine - occupies a different number of ranks in the NCCL group. - - Trainer rank 0 uses ``NCCLWeightTransferEngine.trainer_init`` - in-process (StatelessProcessGroup + PyNcclCommunicator). - """ - if engine_gpu_counts is None: - engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) - - master_address = ray._private.services.get_node_ip_address() - with socket.socket() as sock: - sock.bind(("", 0)) - master_port = sock.getsockname()[1] - world_size = sum(engine_gpu_counts) + 1 # +1 for training rank 0 - - cumulative = [0] - for c in engine_gpu_counts: - cumulative.append(cumulative[-1] + c) - - backend = "hccl" if is_npu() else "nccl" - refs = [ - engine.init_weights_update_group.remote( - master_address=master_address, - master_port=master_port, - rank_offset=cumulative[i] + 1, - world_size=world_size, - group_name=group_name, - backend=backend, - ) - for i, engine in enumerate(rollout_engines) - ] - - if is_npu(): - torch.npu.synchronize() - torch.npu.empty_cache() - device = torch.npu.current_device() - logger.info( - "vLLM in-process weight transfer: addr=%s port=%d world_size=%d device=%d CVD=%s", - master_address, - master_port, - world_size, - device, - os.environ.get("ASCEND_RT_VISIBLE_DEVICES", ""), - ) - # 使用HCCLWeightTransferEngine - from vllm_ascend.distributed.weight_transfer.hccl_engine import HCCLWeightTransferEngine - - model_update_groups = HCCLWeightTransferEngine.trainer_init( - { - "master_address": master_address, - "master_port": master_port, - "world_size": world_size, - } - ) - else: - torch.cuda.synchronize() - torch.cuda.empty_cache() - device = torch.cuda.current_device() - logger.info( - "vLLM in-process weight transfer: addr=%s port=%d world_size=%d device=%d CVD=%s", - master_address, - master_port, - world_size, - device, - os.environ.get("CUDA_VISIBLE_DEVICES", ""), - ) - model_update_groups = NCCLWeightTransferEngine.trainer_init( - { - "master_address": master_address, - "master_port": master_port, - "world_size": world_size, - } - ) - - ray.get(refs) - return model_update_groups - - -def disconnect_rollout_engines_from_distributed( - args: Namespace, - group_name: str, - model_update_groups: Any, - rollout_engines: Sequence[ActorHandle], -) -> None: - """ - Tear down the weight-update NCCL group on the rollout engines. - - ``model_update_groups`` is a vLLM ``PyNcclCommunicator`` returned by - ``NCCLWeightTransferEngine.trainer_init`` (built on a ``StatelessProcessGroup``), - NOT a torch c10d ``ProcessGroup``. It is deliberately not registered in - torch.distributed's global registry, so ``dist.destroy_process_group`` on it - raises ``ValueError: Invalid process group specified`` (see #127 regression). - - We therefore do not tear the trainer-side communicator down here; this matches - the pre-#127 behavior. (Note ``engine.destroy_weights_update_group`` is itself - a no-op on the engine side.) An explicit ``model_update_groups.destroy()`` would - abort the NCCL comm, but that changes long-standing behavior and risks the - CUDA-graph-capture self-deadlock documented in ``PyNcclCommunicator.destroy``; - leave it out of this fix. - """ - refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] - ray.get(refs) - - -def update_weights_from_distributed( - group_name: str, - group: Any, - weight_version: int, - rollout_engines: Sequence[ActorHandle], - converted_named_tensors: Sequence[tuple[str, torch.Tensor]], - *, - packed: bool = False, -) -> list[ObjectRef]: - """ - Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). - - The *group* is a vLLM ``PyNcclCommunicator`` from ``trainer_init`` - in the Megatron trainer process. - """ - refs = [ - engine.update_weights_from_distributed.remote( - names=[name for name, _ in converted_named_tensors], - dtypes=[param.dtype for _, param in converted_named_tensors], - shapes=[param.shape for _, param in converted_named_tensors], - group_name=group_name, - weight_version=str(weight_version), - packed=packed, - ) - for engine in rollout_engines - ] - - named_gpu_iter = ( - (name, (param.data if hasattr(param, "data") else param).contiguous()) - for name, param in converted_named_tensors - ) - if is_npu(): - HCCLWeightTransferEngine.trainer_send_weights( - named_gpu_iter, - HCCLTrainerSendWeightsArgs(group=group, packed=packed), - ) - else: - NCCLWeightTransferEngine.trainer_send_weights( - named_gpu_iter, - NCCLTrainerSendWeightsArgs(group=group, packed=packed), - ) - - return refs - def post_process_weights( restore_weights_before_load: bool, post_process_quantization: bool, rollout_engines: Sequence[ActorHandle], -): - """ - Trigger post-process for int4/fp4 quantization on all rollout engines. - """ +) -> None: ray.get( [ engine.post_process_weights.remote( diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index c1052e4ea..22c7ad8ef 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -1,37 +1,98 @@ -""" -Colocated vLLM weight sync using native IPC transfer engines. -""" - from __future__ import annotations - -import os from argparse import Namespace +from collections import defaultdict from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict +from functools import partial +from typing import Any import ray import torch import torch.distributed as dist from megatron.core import mpu +from ray import ObjectRef from ray.actor import ActorHandle +from tqdm import tqdm -from vime.utils.common import is_npu +from vime.platforms import current_platform +from vime.utils import accelerator from vime.utils.distributed_utils import get_gloo_group +from vime.utils.types import ParamInfo + +from ..dspark.export import export_dspark_model_weights +from ..megatron_to_hf import convert_to_hf +from .common import HfWeightSource, VimeRayWeightSyncClient, create_nccl_trainer +from .expert_routing import configure_expert_routing +from .hf_weight_iterator_direct import HfWeightIteratorDirect +from .update_weight_from_distributed import post_process_weights + + +def _current_gpu_uuid() -> str: + platform = current_platform() + if platform.is_npu: + return platform.weight_transfer.current_device_uuid() + + device_index = torch.cuda.current_device() + props = torch.cuda.get_device_properties(device_index) + return str(props.uuid) + + +def _native_ipc_buffer_size(args: Namespace, param_info_buckets: Sequence[Sequence[ParamInfo]] | None) -> int: + buffer_size = args.update_weight_buffer_size + if not param_info_buckets: + return buffer_size + + tensor_parallel_size = mpu.get_tensor_model_parallel_world_size() + expert_tensor_parallel_size = mpu.get_expert_tensor_parallel_world_size() + for bucket in param_info_buckets: + for info in bucket: + parallel_size = expert_tensor_parallel_size if ".experts." in info.name else tensor_parallel_size + buffer_size = max(buffer_size, info.size * parallel_size) + return buffer_size + + +def _build_packed_ipc_update_info( + named_tensors: Sequence[tuple[str, torch.Tensor]], +) -> tuple[dict[str, Any], torch.Tensor | None]: + if not named_tensors: + return ( + { + "names": [], + "dtype_names": [], + "shapes": [], + "tensor_sizes": [], + "ipc_handles": {}, + }, + None, + ) -from .hf_weight_iterator_base import HfWeightIteratorBase -from .update_weight_from_distributed import ( - connect_rollout_engines_from_distributed, - disconnect_rollout_engines_from_distributed, - post_process_weights, - update_weights_from_distributed, -) + from torch.multiprocessing.reductions import reduce_tensor + from vllm.distributed.weight_transfer.packed_tensor import pack_tensors + + chunk = pack_tensors( + iter(named_tensors), + post_iter_func=lambda item: item[1], + buffer_size_bytes=sum(tensor.numel() * tensor.element_size() for _, tensor in named_tensors), + ) + assert chunk is not None + _, ipc_args = reduce_tensor(chunk.packed_tensor) + gpu_uuid = _current_gpu_uuid() + return ( + { + "names": chunk.names, + "dtype_names": [str(dtype).split(".")[-1] for dtype in chunk.dtypes], + "shapes": chunk.shapes, + "tensor_sizes": chunk.tensor_sizes, + "ipc_handles": {gpu_uuid: ipc_args}, + }, + chunk.packed_tensor, + ) class UpdateWeightFromTensor: """ Update rollout engines from tensor dict: gather TP(GPU NCCL) → convert HF(GPU) → send. - Colocated: build CUDA IPC handles → all_gather_object(Gloo CPU, over the engine + Colocated: build CUDA IPC handles → gather_object(Gloo CPU, over the engine slot ranks) → Ray IPC to engine. Distributed: GPU NCCL broadcast to remote engines. """ @@ -51,24 +112,36 @@ def __init__( self.args = args self.model = model self.weights_getter = weights_getter + self.rank = dist.get_rank() self.model_name = model_name self.quantization_config = quantization_config self.weight_version = 0 self.update_weight_metrics: dict[str, float] = {} - self._hf_weight_iterator = HfWeightIteratorBase.create( + self._hf_weight_iterator = HfWeightIteratorDirect( args=args, model=model, model_name=model_name, quantization_config=quantization_config ) + param_info_buckets = getattr(self._hf_weight_iterator, "megatron_local_param_info_buckets", None) + self._full_param_info_buckets = ( + tuple(tuple(bucket) for bucket in param_info_buckets) if param_info_buckets is not None else None + ) + self._non_expert_param_info_buckets: list[list[ParamInfo]] | None = None + draft_weights_getter = ( + partial( + export_dspark_model_weights, + self.model, + use_policy_embedding=not self.args.dspark_pretrained_model, + ) + if self.args.dspark_enabled + else None + ) + self._source = HfWeightSource(self._hf_weight_iterator, self.weights_getter, draft_weights_getter) - self._model_update_groups = None - # vLLM #39212 IPC transfer-engine init runs once per set of colocated engines. - self._ipc_initialized = False - # vLLM IPC handle payloads are pickled on the Ray/HTTP bridge. - os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - - # ------------------------------------------------------------------ - # connect / disconnect - # ------------------------------------------------------------------ + self._ipc_gather_group = None + self._ipc_gather_src = None + self._ipc_engine = None + self._expert_transfer_plan = [] + self._native_trainers = [] def connect_rollout_engines( self, @@ -76,12 +149,15 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, + engine_parallel_configs: Sequence[Mapping[str, Any]] | None = None, ) -> None: - """ - Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL - for distributed. Map ranks to colocated IPC engines. - """ - self.rollout_engines = rollout_engines + del rollout_engine_lock + for trainer in self._native_trainers: + trainer.shutdown() + self._all_rollout_engines = list(rollout_engines) + self.rollout_engines = [] + self._ipc_engine = None + self._native_trainers = [] if engine_gpu_counts is None: engine_gpu_counts = [self.args.rollout_num_gpus_per_engine] * len(rollout_engines) @@ -101,46 +177,174 @@ def connect_rollout_engines( break colocate_engine_nums += 1 - self.use_distribute = len(rollout_engines) > colocate_engine_nums + self.rollout_engines = list(rollout_engines[:colocate_engine_nums]) + distributed_rollout_engines = list(rollout_engines[colocate_engine_nums:]) + use_distribute = bool(distributed_rollout_engines) + colocate_gpu_offsets = engine_gpu_offsets[:colocate_engine_nums] + colocate_gpu_counts = engine_gpu_counts[:colocate_engine_nums] + colocate_parallel_configs = ( + engine_parallel_configs[:colocate_engine_nums] if engine_parallel_configs is not None else None + ) - if self.use_distribute: - self.rollout_engines = rollout_engines[:colocate_engine_nums] - self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:] - distributed_gpu_counts = engine_gpu_counts[colocate_engine_nums:] - self._is_distributed_src_rank = ( - mpu.get_data_parallel_rank(with_context_parallel=True) == 0 - and mpu.get_tensor_model_parallel_rank() == 0 - and mpu.get_pipeline_model_parallel_rank() == 0 - ) - self._group_name = "vime" - if self._is_distributed_src_rank: - if self._model_update_groups is not None: - disconnect_rollout_engines_from_distributed( - self.args, self._group_name, self._model_update_groups, self.distributed_rollout_engines - ) - self._model_update_groups = connect_rollout_engines_from_distributed( - self.args, - self._group_name, - self.distributed_rollout_engines, - engine_gpu_counts=distributed_gpu_counts, + self._non_expert_param_info_buckets, self._expert_transfer_plan = configure_expert_routing( + args=self.args, + full_param_info_buckets=self._full_param_info_buckets, + get_local_weight_names=self.weights_getter, + engine_gpu_counts=colocate_gpu_counts, + engine_gpu_offsets=colocate_gpu_offsets, + engine_parallel_configs=colocate_parallel_configs, + use_distribute=use_distribute, + ) + + if not self._expert_transfer_plan: + if self.rollout_engines: + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + client = VimeRayWeightSyncClient(self.rollout_engines, lambda: self.weight_version) + trainer = WeightTransferTrainerFactory.trainer_init( + current_platform().weight_transfer.trainer_init_info( + colocate=True, + rank=dist.get_rank(), + packed=True, + packed_buffer_size_bytes=_native_ipc_buffer_size( + self.args, + self._full_param_info_buckets, + ), + ), + client=client, + source=self._source, ) + self._native_trainers.append(trainer) + if distributed_rollout_engines: + distributed_gpu_counts = engine_gpu_counts[colocate_engine_nums:] + client = VimeRayWeightSyncClient( + distributed_rollout_engines, + lambda: self.weight_version, + distributed_gpu_counts, + ) + trainer = create_nccl_trainer( + client, + self._source, + distributed_gpu_counts, + ) + self._native_trainers.append(trainer) + + return - # vLLM #39212: one-time IPC transfer-engine init on each colocated engine. - if dist.get_rank() == 0 and self.rollout_engines and not self._ipc_initialized: - ray.get([engine.init_weight_transfer_engine.remote({"init_info": {}}) for engine in self.rollout_engines]) - self._ipc_initialized = True + # Rank-local expert routing is the one case the generic IPC API cannot + # express: each rollout EP rank receives a different expert subset. + if self._ipc_gather_group is None: + for index in range(colocate_engine_nums): + group_ranks = list( + range( + colocate_gpu_offsets[index], + colocate_gpu_offsets[index] + colocate_gpu_counts[index], + ) + ) + new_group = dist.new_group(ranks=group_ranks, backend="gloo") + if dist.get_rank() in group_ranks: + self._ipc_gather_group = new_group + self._ipc_gather_src = colocate_gpu_offsets[index] + + for index, engine in enumerate(self.rollout_engines): + start = colocate_gpu_offsets[index] + if start <= dist.get_rank() < start + colocate_gpu_counts[index]: + self._ipc_engine = engine + + if dist.get_rank() == 0: + ray.get( + [ + engine.init_weight_transfer_engine.remote({"init_info": {"packed": True}}) + for engine in self.rollout_engines + ] + ) def pop_metrics(self) -> dict[str, float]: - """ - Return and clear ``update_weight_metrics``. Empty under colocate today; - kept symmetric with UpdateWeightFromDistributed so the actor can drain unconditionally. - """ out, self.update_weight_metrics = self.update_weight_metrics, {} return out - # ------------------------------------------------------------------ - # weight update - # ------------------------------------------------------------------ + def _prepare_expert_weight_batch( + self, + transfers: Sequence[Any], + megatron_local_weights: Mapping[str, torch.Tensor], + staging_buffers: dict[tuple[torch.dtype, tuple[int, ...]], list[torch.Tensor]], + ) -> list[tuple[str, torch.Tensor]]: + local_params = [] + p2p_ops = [] + buffer_offsets: dict[tuple[torch.dtype, tuple[int, ...]], int] = defaultdict(int) + for transfer in transfers: + for expert_param in transfer.params: + info = expert_param.info + if self.rank != transfer.source_rank and self.rank not in transfer.target_ranks: + continue + key = (info.dtype, tuple(info.shape)) + pool = staging_buffers.setdefault(key, []) + offset = buffer_offsets[key] + buffer_offsets[key] = offset + 1 + if offset == len(pool): + pool.append(torch.empty(info.shape, dtype=info.dtype, device=accelerator.device())) + tensor = pool[offset] + if self.rank == transfer.source_rank: + source = megatron_local_weights[info.name] + if source.shape != info.shape or source.dtype != info.dtype: + raise ValueError(f"expert metadata changed for {info.name}") + tensor.copy_(source, non_blocking=True) + p2p_ops.extend( + dist.P2POp(dist.isend, tensor, target_rank) + for target_rank in transfer.target_ranks + if target_rank != self.rank + ) + if self.rank in expert_param.target_ranks: + local_params.append((expert_param, tensor)) + else: + p2p_ops.append(dist.P2POp(dist.irecv, tensor, transfer.source_rank)) + local_params.append((expert_param, tensor)) + + for request in dist.batch_isend_irecv(p2p_ops) if p2p_ops else (): + request.wait() + + hf_named_tensors = [] + for expert_param, tensor in local_params: + hf_named_tensors.extend( + convert_to_hf( + self.args, + self.model_name, + expert_param.info.name, + tensor, + self.quantization_config, + ) + ) + return hf_named_tensors + + def _update_expert_weights( + self, + megatron_local_weights: Mapping[str, torch.Tensor], + ) -> None: + dist.barrier(group=get_gloo_group()) + # Initialize WORLD on all ranks before subset batched P2P. + dist.barrier() + # Reuse staging across layers instead of fragmenting the CUDA allocator. + staging_buffers: dict[tuple[torch.dtype, tuple[int, ...]], list[torch.Tensor]] = {} + for transfer_group in tqdm( + self._expert_transfer_plan, + disable=self.rank != 0, + desc="Update expert weights", + ): + for transfer_batch in transfer_group: + hf_named_tensors = self._prepare_expert_weight_batch( + transfer_batch, + megatron_local_weights, + staging_buffers, + ) + refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) + ray.get(refs) + dist.barrier(group=get_gloo_group()) + accelerator.synchronize() + del refs, long_lived_tensors, hf_named_tensors + accelerator.ipc_collect() + accelerator.empty_cache() + del staging_buffers + accelerator.empty_cache() @torch.no_grad() def update_weights(self) -> None: @@ -148,291 +352,118 @@ def update_weights(self) -> None: version++, flush caches, process buckets. Progress on rank 0. """ self.weight_version += 1 - - rank = dist.get_rank() - if rank == 0: - ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) - ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + if self.rank == 0: + ray.get([engine.pause_generation.remote() for engine in self._all_rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self._all_rollout_engines]) if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: post_process_weights( restore_weights_before_load=True, post_process_quantization=False, - rollout_engines=self.rollout_engines, + rollout_engines=self._all_rollout_engines, ) dist.barrier(group=get_gloo_group()) - # Enter the native vLLM weight-update state machine on every colocated engine. - if rank == 0: - ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in self.rollout_engines]) - dist.barrier(group=get_gloo_group()) - - megatron_local_weights = self.weights_getter() - - for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): - refs = self._send_hf_params(hf_named_tensors) - ray.get(refs) - # Free chunk tensors so the caching allocator can reuse the blocks. - del hf_named_tensors - if is_npu(): - torch.npu.synchronize() - else: - torch.cuda.ipc_collect() - - dist.barrier(group=get_gloo_group()) - # After the barrier all engines have returned, so every rank's last-chunk - # IPC handles are now released by the consumers. Clean them up. - if is_npu(): - torch.npu.synchronize() + if self._native_trainers: + for trainer in self._native_trainers: + trainer.client.draft = False + trainer.send_weights() + update_draft = self.args.dspark_enabled or ( + self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp" + ) + if update_draft: + self._source.draft = self.args.dspark_enabled + for trainer in self._native_trainers: + trainer.client.draft = True + trainer.send_weights() + trainer.client.draft = False + self._source.draft = False else: - torch.cuda.ipc_collect() + megatron_local_weights = self.weights_getter() + self._update_rollout_weights(megatron_local_weights, draft=False) - # Exit the native vLLM weight-update state machine. - if rank == 0: - ray.get([engine.finish_weight_update.remote() for engine in self.rollout_engines]) - dist.barrier(group=get_gloo_group()) + if self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp": + self._update_rollout_weights(megatron_local_weights, draft=True) # int4/fp4 post_process - if rank == 0: + if self.rank == 0: if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: post_process_weights( restore_weights_before_load=False, post_process_quantization=True, - rollout_engines=self.rollout_engines, + rollout_engines=self._all_rollout_engines, ) - ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.continue_generation.remote() for engine in self._all_rollout_engines]) dist.barrier(group=get_gloo_group()) - def _send_hf_params(self, hf_named_tensors) -> list[object]: - all_refs: list[object] = [] + def _update_rollout_weights(self, megatron_local_weights, *, draft: bool) -> None: + if self._ipc_engine is not None and self.rank == self._ipc_gather_src: + method = self._ipc_engine.start_draft_weight_update if draft else self._ipc_engine.start_weight_update + ray.get(method.remote()) + dist.barrier(group=get_gloo_group()) - _send_to_colocated_engine( - hf_named_tensors, - rollout_engines=self.rollout_engines, - weight_version=self.weight_version, - ) + self._send_weight_chunks(megatron_local_weights) + dist.barrier(group=get_gloo_group()) + accelerator.ipc_collect() + accelerator.empty_cache() - if self.use_distribute and self._is_distributed_src_rank: - refs_distributed = update_weights_from_distributed( - self._group_name, - self._model_update_groups, - self.weight_version, - self.distributed_rollout_engines, - hf_named_tensors, - packed=False, - ) - if refs_distributed: - all_refs.extend(refs_distributed) + if self._ipc_engine is not None and self.rank == self._ipc_gather_src: + ray.get(self._ipc_engine.finish_weight_update.remote(weight_version=str(self.weight_version))) + dist.barrier(group=get_gloo_group()) - return all_refs + def _send_weight_chunks(self, megatron_local_weights) -> None: + param_info_buckets = ( + self._non_expert_param_info_buckets if self._expert_transfer_plan else self._full_param_info_buckets + ) + for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks( + megatron_local_weights, + param_info_buckets=param_info_buckets, + ): + refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) + ray.get(refs) + del refs, long_lived_tensors, hf_named_tensors + accelerator.ipc_collect() + accelerator.empty_cache() + if self._expert_transfer_plan: + self._update_expert_weights(megatron_local_weights) + + def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: + return _send_to_colocated_engine( + hf_named_tensors, + ipc_engine=self._ipc_engine, + ipc_gather_src=self._ipc_gather_src, + ipc_gather_group=self._ipc_gather_group, + ) def _send_to_colocated_engine( hf_named_tensors: list[tuple[str, torch.Tensor]], *, - rollout_engines: Sequence[ActorHandle], - weight_version: int, -) -> None: - if not rollout_engines: - return - - def send_to_vllm(update_info) -> None: - request = {"update_info": asdict(update_info)} - ray.get( - [engine.update_weights.remote(request, weight_version=str(weight_version)) for engine in rollout_engines] - ) - - if is_npu(): - from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import ( - NPUIPCTrainerSendWeightsArgs, - NPUIPCWeightTransferEngine, - ) - - trainer_args = NPUIPCTrainerSendWeightsArgs(send_mode=send_to_vllm, packed=False) - NPUIPCWeightTransferEngine.trainer_send_weights(iter(hf_named_tensors), trainer_args) - else: - from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerSendWeightsArgs, IPCWeightTransferEngine - - trainer_args = IPCTrainerSendWeightsArgs(send_mode=send_to_vllm, packed=False) - IPCWeightTransferEngine.trainer_send_weights(iter(hf_named_tensors), trainer_args) - - -# --------------------------------------------------------------------------- -# vLLM worker extension (loaded by ``--worker-extension-cls``) -# --------------------------------------------------------------------------- - - -class _VLLMHijack: - """vLLM worker extension helpers. - - On NPU: - - Patches NPUWorker.load_model and NPUWorker.start_weight_update to fix - MoE weight_loader missing on EP (a vLLM bug where w13_weight/w2_weight - params lack weight_loader attr when EP is enabled). - - Patches ApplyRotaryEmb.__init__ to skip flash_attn import - (Megatron/NPU backends introduce flash_attn as a dummy module, - but vllm_ascend does not use it). - """ - - @staticmethod - def _patch_npu_worker() -> None: - from vllm_ascend.worker.worker import NPUWorker - - if getattr(NPUWorker, "_npu_worker_patched", False): - return - - _VLLMHijack._patch_one_worker(NPUWorker) - NPUWorker._npu_worker_patched = True - - @staticmethod - def _patch_a3_moe_alltoall_expert_ids() -> None: - """Restore the ALLTOALL expert-ID template after colocated memory reuse.""" - from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type - - if get_ascend_device_type() != AscendDeviceType.A3: - return - - from vllm_ascend.ops.fused_moe.token_dispatcher import TokenDispatcherWithAll2AllV - - if getattr(TokenDispatcherWithAll2AllV, "_vime_expert_ids_patched", False): - return - - original_dispatch_preprocess = TokenDispatcherWithAll2AllV._dispatch_preprocess - TokenDispatcherWithAll2AllV._vime_expert_ids_generation = 0 - - def _patched_dispatch_preprocess(self, hidden_states, topk_ids): - generation = TokenDispatcherWithAll2AllV._vime_expert_ids_generation - if self.num_local_experts > 1 and getattr(self, "_vime_seen_expert_ids_generation", -1) != generation: - expert_ids = self.expert_ids_per_ep_rank - self.expert_ids_per_ep_rank = torch.arange( - self.num_experts, - device=expert_ids.device, - dtype=expert_ids.dtype, - ).remainder(self.num_local_experts) - self._vime_seen_expert_ids_generation = generation - return original_dispatch_preprocess(self, hidden_states, topk_ids) - - TokenDispatcherWithAll2AllV._dispatch_preprocess = _patched_dispatch_preprocess - TokenDispatcherWithAll2AllV._vime_expert_ids_patched = True - - @staticmethod - def _invalidate_moe_alltoall_expert_ids() -> None: - try: - from vllm_ascend.ops.fused_moe.token_dispatcher import TokenDispatcherWithAll2AllV - except ImportError: - return - - if getattr(TokenDispatcherWithAll2AllV, "_vime_expert_ids_patched", False): - TokenDispatcherWithAll2AllV._vime_expert_ids_generation += 1 - - @staticmethod - def _patch_one_worker(worker_cls: type) -> None: - import inspect - - _orig_load_model = worker_cls.load_model - _orig_start_weight_update = worker_cls.start_weight_update - _orig_wake_up = worker_cls.wake_up - has_dummy_kw = "load_dummy_weights" in inspect.signature(_orig_load_model).parameters - - if has_dummy_kw: - - def _patched_load_model(self, *, load_dummy_weights: bool = False, _orig=_orig_load_model) -> None: - _orig(self, load_dummy_weights=load_dummy_weights) - _VLLMHijack.patch_moe_weight_loader(self.model_runner.model) - - else: - - def _patched_load_model(self, _orig=_orig_load_model) -> None: - _orig(self) - _VLLMHijack.patch_moe_weight_loader(self.model_runner.model) - - def _patched_start_weight_update( - self, is_checkpoint_format: bool = True, _orig=_orig_start_weight_update - ) -> None: - _VLLMHijack.patch_moe_weight_loader(self.model_runner.model) - _orig(self, is_checkpoint_format=is_checkpoint_format) - _VLLMHijack._invalidate_moe_alltoall_expert_ids() - - def _patched_wake_up(self, tags=None, _orig=_orig_wake_up) -> None: - quant_config = self.vllm_config.quant_config - if quant_config is not None: - _orig(self, tags=tags) - _VLLMHijack._invalidate_moe_alltoall_expert_ids() - return - - # vllm-ascend transposes unquantized w13_weight/w2_weight in - # wake_up(). Keep the native allocator and buffer restoration, but - # skip that branch: layerwise reload owns the final runtime layout. - self.vllm_config.quant_config = object() - try: - _orig(self, tags=tags) - finally: - self.vllm_config.quant_config = quant_config - _VLLMHijack._invalidate_moe_alltoall_expert_ids() - - worker_cls.load_model = _patched_load_model # type: ignore[attr-defined] - worker_cls.start_weight_update = _patched_start_weight_update # type: ignore[attr-defined] - worker_cls.wake_up = _patched_wake_up # type: ignore[attr-defined] - - @staticmethod - def patch_moe_weight_loader(model: torch.nn.Module) -> None: - inner_model = getattr(model, "model", None) or getattr(model, "language_model", None) - if inner_model is None: - return - if not hasattr(inner_model, "layers"): - inner_model = getattr(inner_model, "model", None) - if inner_model is None or not hasattr(inner_model, "layers"): - return - - for layer in inner_model.layers: - mlp = getattr(layer, "mlp", None) or getattr(layer, "block_sparse_moe", None) - if mlp is None: - continue - experts = getattr(mlp, "experts", None) - if experts is None or not hasattr(experts, "weight_loader"): - continue - for name, param in mlp.named_parameters(): - if "w13_weight" in name or "w2_weight" in name: - if not hasattr(param, "weight_loader"): - param.weight_loader = experts.weight_loader # type: ignore[attr-defined] - - @staticmethod - def _patch_npu_rotary_emb() -> None: - from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb - - if getattr(ApplyRotaryEmb, "_npu_rotary_patched", False): - return - - def _npu_rotary_emb_init( - self, - enforce_enable: bool = False, - is_neox_style: bool = True, - enable_fp32_compute: bool = False, - ) -> None: - super(ApplyRotaryEmb, self).__init__(enforce_enable=enforce_enable) - self.is_neox_style = is_neox_style - self.enable_fp32_compute = enable_fp32_compute - self.apply_rotary_emb_flash_attn = None - - ApplyRotaryEmb.__init__ = _npu_rotary_emb_init # type: ignore[attr-defined] - ApplyRotaryEmb._npu_rotary_patched = True - - -class vLLMColocateWorkerExtension: - """vLLM ``--worker-extension-cls`` entry for colocated rollout workers.""" - - def __new__(cls, **kwargs): - if is_npu(): - _VLLMHijack._patch_a3_moe_alltoall_expert_ids() - _VLLMHijack._patch_npu_worker() - _VLLMHijack._patch_npu_rotary_emb() - return super().__new__(cls) - - -class vLLMWorkerExtension: - """vLLM ``--worker-extension-cls`` entry for general bugfix.""" - - def __new__(cls, **kwargs): - if is_npu(): - _VLLMHijack._patch_npu_worker() - _VLLMHijack._patch_npu_rotary_emb() - return super().__new__(cls) + ipc_engine, + ipc_gather_src, + ipc_gather_group, +) -> tuple[list[ObjectRef], Any]: + # Placeholder ranks (GPU slots reserved but no engine) have no gather group. + # gather_object is only collective among group members, so we skip entirely. + if ipc_gather_group is None: + return [], None + + local_info, weight_ref = _build_packed_ipc_update_info(hf_named_tensors) + + slot_size = dist.get_world_size(ipc_gather_group) + if slot_size <= 1: + if not local_info["names"]: + return [], weight_ref + ref = ipc_engine.update_weights.remote(local_info) + return [ref], weight_ref + + gathered_infos = [None] * slot_size if dist.get_rank() == ipc_gather_src else None + dist.gather_object(local_info, object_gather_list=gathered_infos, dst=ipc_gather_src, group=ipc_gather_group) + + refs = [] + if dist.get_rank() == ipc_gather_src: + if any(info is None for info in gathered_infos): + raise RuntimeError(f"Missing IPC payloads in slot {ipc_gather_src}; got {gathered_infos!r}") + if any(info["names"] for info in gathered_infos): + refs.append(ipc_engine.update_weights.remote(gathered_infos)) + + return refs, weight_ref diff --git a/vime/backends/megatron_utils/vllm.py b/vime/backends/megatron_utils/vllm.py new file mode 100644 index 000000000..0224de600 --- /dev/null +++ b/vime/backends/megatron_utils/vllm.py @@ -0,0 +1,52 @@ +"""vLLM FP8 helpers used by Megatron weight conversion.""" + +from math import ceil + +import torch +from vllm.utils.deep_gemm import ( + get_mn_major_tma_aligned_packed_ue8m0_tensor, + get_tma_aligned_size, + is_deep_gemm_e8m0_used, + per_block_cast_to_fp8, +) + + +def should_deepgemm_weight_requant_ue8m0(weight_block_size) -> bool: + return weight_block_size is not None and is_deep_gemm_e8m0_used() + + +def quant_weight_ue8m0( + weight_dequant: torch.Tensor, + weight_block_size: list[int], +): + assert weight_block_size == [128, 128] + assert weight_dequant.dtype == torch.bfloat16, f"{weight_dequant.dtype=} {weight_dequant.shape=}" + *batch_dims, n, k = weight_dequant.shape + flat = weight_dequant.view(-1, k) + out_w_flat, out_s_flat = per_block_cast_to_fp8(flat, block_size=[128, 128], use_ue8m0=True) + out_w = out_w_flat.view(*batch_dims, n, k) + out_s = out_s_flat.view( + *batch_dims, + ceil(n / weight_block_size[0]), + ceil(k / weight_block_size[1]), + ) + return out_w, out_s + + +def transform_scale_ue8m0(sf: torch.Tensor, mn: int): + sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128) + sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) + if sf.shape[-1] == 1: + aligned_mn = get_tma_aligned_size(sf.shape[-2], sf.element_size()) + if sf.stride(-1) != aligned_mn: + new_stride = list(sf.stride()) + new_stride[-1] = aligned_mn + sf = sf.as_strided(sf.shape, tuple(new_stride)) + return sf + + +__all__ = [ + "quant_weight_ue8m0", + "transform_scale_ue8m0", + "should_deepgemm_weight_requant_ue8m0", +] diff --git a/vime/backends/vllm_utils/__init__.py b/vime/backends/vllm_utils/__init__.py index e69de29bb..f372cb001 100644 --- a/vime/backends/vllm_utils/__init__.py +++ b/vime/backends/vllm_utils/__init__.py @@ -0,0 +1,5 @@ +from vime.utils import accelerator + +# Finalize the backend before importing any vLLM submodule. In a MUSA +# runtime this loads musa_patch only after MUSA wins backend selection. +accelerator.initialize_accelerator() diff --git a/vime/backends/vllm_utils/arguments.py b/vime/backends/vllm_utils/arguments.py index f6a6aa348..ccb1c7863 100644 --- a/vime/backends/vllm_utils/arguments.py +++ b/vime/backends/vllm_utils/arguments.py @@ -1,225 +1,112 @@ -"""vLLM rollout backend argument definitions. - -Wholesale-imports ``AsyncEngineArgs.add_cli_args`` prefixed ``--vllm-``/``vllm_``, -adds vime orchestration extras, and provides ``get_vllm_cli_action_table`` for -subprocess CLI forwarding (vLLM is launched as ``vllm serve``). -""" - import argparse -import sys +import logging from vllm.engine.arg_utils import AsyncEngineArgs from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_router.launch_router import RouterArgs from vime.utils.http_utils import _wrap_ipv6 - -def _detect_user_provided_dests(parser, argv: list[str]) -> tuple[set[str], dict[str, str]]: - """Return (user_provided, raw_values) extracted from ``argv``. - - ``user_provided``: dests explicitly named on the CLI — used to distinguish "user - passed a value equaling the default" from "user accepted the parsed default". - - ``raw_values``: literal CLI strings per dest — used to forward dataclass-backed - flags (e.g. ``--vllm-compilation-config``) verbatim instead of re-serializing - the parsed runtime object. - """ - flag_to_dest: dict[str, str] = {} - for action in parser._actions: - for flag in action.option_strings: - flag_to_dest[flag] = action.dest - user: set[str] = set() - raw: dict[str, str] = {} - i = 0 - while i < len(argv): - token = argv[i] - if "=" in token and token.startswith("--"): - head, raw_val = token.split("=", 1) - dest = flag_to_dest.get(head) - if dest is not None: - user.add(dest) - raw[dest] = raw_val - i += 1 - continue - dest = flag_to_dest.get(token) - if dest is not None: - user.add(dest) - if i + 1 < len(argv) and not argv[i + 1].startswith("--"): - raw[dest] = argv[i + 1] - i += 2 - continue - i += 1 - return user, raw - - -# Dests orchestrator-owned or non-applicable to subprocess `vllm serve` mode. -SKIPPED_DESTS = [ - "model", - "served_model_name", - "config", - "tokenizer", - "tokenizer_mode", - "tokenizer_revision", - "trust_remote_code", - "seed", - "dtype", - # TP is orchestrator-owned; PP/DP remain user-controllable and auto-forward. - "tensor_parallel_size", - "nnodes", - "node_rank", - "master_addr", - "master_port", - "data_parallel_backend", - "distributed_executor_backend", - "port", - "host", - "enable_return_routed_experts", -] +logger = logging.getLogger(__name__) def add_vllm_router_arguments(parser): - """vime's vllm-router endpoint flags (host/port are orchestrator-owned, excluded from RouterArgs CLI).""" parser.add_argument( "--vllm-router-ip", type=str, default=None, - help="IP address of the vllm router (where vime connects to send rollout requests).", + help="IP address of the vllm router", ) parser.add_argument( "--vllm-router-port", type=int, default=None, - help="Port of the vllm router.", + help="Port of the vllm router", ) - # Bare --router-* namespace: this is a real vllm-router knob (RouterArgs field); - # host/port use --vllm-router-* because RouterArgs excludes them (exclude_host_port=True). parser.add_argument( - "--router-request-timeout-secs", + "--vllm-router-request-timeout-secs", type=int, default=14400, - help="Timeout (seconds) for HTTP requests vime makes to the vllm router.", - ) - # dest=router_policy (not vllm_router_policy): flows into RouterArgs.from_cli_args and - # is read by vllm_rollout.generate to decide whether to send x-session-id headers. - parser.add_argument( - "--vllm-router-policy", - type=str, - default="consistent_hash", - dest="router_policy", - choices=["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash"], - help=( - "vllm-router load-balancing policy. Defaults to 'consistent_hash' for " - "session-affinity routing replay via the x-session-id header." - ), + help="Timeout for requests to the vllm router in seconds", ) + RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True) return parser -def _make_add_argument_wrapper(target_add_argument): - """Return a wrapper that skips dests in SKIPPED_DESTS and prefixes flags/dest with ``vllm-``/``vllm_``.""" - - def wrapper(*name_or_flags, **kwargs): - # determine canonical dest for skip check - canonical = kwargs.get("dest") - if canonical is None: - for s in name_or_flags: - if isinstance(s, str) and s.startswith("--"): - canonical = s[2:].replace("-", "_") - break - if canonical in SKIPPED_DESTS: - return None - - # prefix flags - new_flags = [] - for s in name_or_flags: - if isinstance(s, str) and s.startswith("-"): - new_flags.append(f"--vllm-{s.lstrip('-')}") - else: - new_flags.append(s) - - # prefix dest - new_kwargs = kwargs.copy() - if "dest" in new_kwargs and isinstance(new_kwargs["dest"], str): - if not new_kwargs["dest"].startswith("vllm_"): - new_kwargs["dest"] = f"vllm_{new_kwargs['dest']}" - - return target_add_argument(*new_flags, **new_kwargs) - - return wrapper - - def add_vllm_arguments(parser): - """Register --vllm-* flags into parser. - - Wholesale-imports ``AsyncEngineArgs.add_cli_args`` via a monkey-patched - ``parser.add_argument`` / ``parser.add_argument_group`` wrapper that prefixes - every flag with ``--vllm-`` and every dest with ``vllm_``, skipping dests in - ``SKIPPED_DESTS``. Both are patched because vLLM creates argument groups. - Pass a ``FlexibleArgumentParser`` so vLLM's ``deprecated`` kwarg is handled - natively on Python 3.12. - """ parser = add_vllm_router_arguments(parser) - parser.add_argument( - "--vllm-server-concurrency", - type=int, - default=512, - help="Max concurrent inference requests sent to each vLLM server worker.", - ) + parser.set_defaults(router_balance_abs_threshold=10, router_balance_rel_threshold=1.2) + parser.add_argument("--vllm-server-concurrency", type=int, default=512) parser.add_argument( "--vllm-enable-deterministic-inference", action="store_true", default=False, help=( "Make rollout sampling deterministic. Forwards a per-sample ``seed`` " - "(derived from ``--rollout-seed`` and the sample's index in the group) " - "AND exports ``VLLM_BATCH_INVARIANT=1`` to the vLLM subprocess so attention " - "/ comm / MM kernels pick batch-invariant variants. Both are required for " - "true determinism — seed alone does not control kernel selection." - ), - ) - parser.add_argument( - "--vllm-tool-call-parser", - dest="vllm_tool_call_parser", - type=str, - default=None, - help="vLLM tool-call parser name for agent output parsing (e.g. qwen3_coder).", - ) - _vllm_packed = parser.add_mutually_exclusive_group() - _vllm_packed.add_argument( - "--vllm-weight-sync-packed", - dest="vllm_weight_sync_packed", - action="store_true", - help=( - "Use one-shot packed weight transfer for dense models (no MoE experts). " - "Automatically disabled for MoE or compressed-tensors quantization." + "AND exports ``VLLM_BATCH_INVARIANT=1`` to the vLLM subprocess." ), ) - _vllm_packed.add_argument( - "--no-vllm-weight-sync-packed", - dest="vllm_weight_sync_packed", - action="store_false", - help="Disable packed sync; send weights per bucket via in-process NCCL (non-packed mode).", - ) + # Monkey-patch parser to prefix all engine flags with --vllm- / vllm_ + old_add_argument = parser.add_argument + old_add_argument_group = parser.add_argument_group + + skipped_args = [ + "model", + "config", + "trust_remote_code", + "seed", + "tensor_parallel_size", + "nnodes", + "node_rank", + "master_addr", + "master_port", + "data_parallel_backend", + "distributed_executor_backend", + "port", + "host", + "enable_return_routed_experts", + ] + + def _wrap_add_argument(target_add_argument): + def wrapper(*name_or_flags, **kwargs): + canonical = kwargs.get("dest") + if canonical is None: + for s in name_or_flags: + if isinstance(s, str) and s.startswith("--"): + canonical = s[2:].replace("-", "_") + break + if canonical in skipped_args: + return None + + new_flags = [] + for s in name_or_flags: + if isinstance(s, str) and s.startswith("-"): + new_flags.append(f"--vllm-{s.lstrip('-')}") + else: + new_flags.append(s) + + new_kwargs = kwargs.copy() + if "dest" in new_kwargs and isinstance(new_kwargs["dest"], str): + if not new_kwargs["dest"].startswith("vllm_"): + new_kwargs["dest"] = f"vllm_{new_kwargs['dest']}" - parser.set_defaults(vllm_weight_sync_packed=True) + return target_add_argument(*new_flags, **new_kwargs) - old_parser_add_argument = parser.add_argument - old_parser_add_argument_group = parser.add_argument_group + return wrapper def patched_add_argument_group(*g_args, **g_kwargs): - group = old_parser_add_argument_group(*g_args, **g_kwargs) - group.add_argument = _make_add_argument_wrapper(group.add_argument) + group = old_add_argument_group(*g_args, **g_kwargs) + group.add_argument = _wrap_add_argument(group.add_argument) return group - parser.add_argument = _make_add_argument_wrapper(old_parser_add_argument) + parser.add_argument = _wrap_add_argument(old_add_argument) parser.add_argument_group = patched_add_argument_group AsyncEngineArgs.add_cli_args(parser) - parser.add_argument = old_parser_add_argument - parser.add_argument_group = old_parser_add_argument_group + from vllm.entrypoints.launchers.cli_args import FrontendArgs - # Deliberately no set_defaults for vllm flags: argparse.set_defaults mutates action.default, - # which would make _forward_vllm_cli_args skip forwarding values that equal the default. - # vime-preferred defaults are applied explicitly in vllm_engine.launch_server_process. + FrontendArgs.add_cli_args(parser) + parser.add_argument = old_add_argument + parser.add_argument_group = old_add_argument_group # PD disaggregation / multi-group config parser.add_argument( @@ -235,7 +122,6 @@ def patched_add_argument_group(*g_args, **g_kwargs): dest="vllm_config", help=( "Path to a YAML config file for fine-grained vLLM rollout engine deployment. " - "Enables multi-model serving, PD disaggregation, and heterogeneous server groups. " "Mutually exclusive with --prefill-num-servers and --rollout-external." ), ) @@ -244,10 +130,12 @@ def patched_add_argument_group(*g_args, **g_kwargs): def validate_args(args): - """vLLM-specific validation.""" args.vllm_dp_size = args.vllm_data_parallel_size args.vllm_pp_size = args.vllm_pipeline_parallel_size + if getattr(args, "rollout_top_p", 1.0) != 1.0 and getattr(args, "rollout_top_k", -1) <= 0: + raise ValueError("vLLM top-p sampling replay requires --rollout-top-k > 0.") + if getattr(args, "vllm_router_ip", None): args.vllm_router_ip = _wrap_ipv6(args.vllm_router_ip) @@ -264,119 +152,26 @@ def validate_args(args): ), "vllm_config and prefill_num_servers are mutually exclusive. Use server_groups in the YAML config instead." -def _sanitize_non_primitive_defaults(args, user_provided: set[str]) -> None: - """Replace non-primitive default values with None. - - vllm's ``AsyncEngineArgs.add_cli_args`` registers parameters whose defaults - are complex dataclass/config instances (e.g. ``EPLBConfig()``, - ``CompilationConfig()``). When ``parse_known_args`` resolves those defaults - they become live Python objects on the returned Namespace. - - The main ``parse_args`` in ``slime/utils/arguments.py`` later merges - *every* attribute from this Namespace into the global ``args`` via - ``setattr``. Downstream code (megatron-bridge, model_provider, …) may - iterate ``vars(args)`` and choke on unrecognised types. - - This function walks the parsed Namespace and replaces any non-primitive - *default* value (i.e. a dest the user did NOT explicitly provide) with - ``None``. User-provided values are intentionally left intact — they are - forwarded to the vllm subprocess via ``_vllm_raw_values`` and do not need - the parsed object representation in the main args namespace. - """ - _PRIMITIVE = (str, int, float, bool, type(None)) - - for key, value in list(vars(args).items()): - if key.startswith("_"): - continue - if key in user_provided: - continue - if isinstance(value, _PRIMITIVE): - continue - if isinstance(value, (list, tuple)): - if all(isinstance(v, _PRIMITIVE) for v in value): - continue - setattr(args, key, None) - - def vllm_parse_args(): """Parse vLLM flags via an independent ArgumentParser + parse_known_args. - Returns a Namespace with all attrs prefixed ``vllm_``, plus: - - ``_vllm_user_provided``: set of dests the user named on argv - - ``_vllm_raw_values``: literal CLI strings per dest (for verbatim forwarding - of dataclass-backed flags like ``--vllm-compilation-config``) + Returns a Namespace with all attrs prefixed ``vllm_``. """ parser = FlexibleArgumentParser(add_help=False) add_vllm_arguments(parser) - # Compute default vllm_tensor_parallel_size from CLI args + # Compute default vllm_tensor_parallel_size from CLI args. temp_parser = argparse.ArgumentParser(add_help=False) temp_parser.add_argument("--rollout-num-gpus-per-engine", type=int, default=1) temp_parser.add_argument("--vllm-pipeline-parallel-size", type=int, default=1) + temp_parser.add_argument("--vllm-prefill-context-parallel-size", type=int, default=1) + temp_parser.add_argument("--vllm-data-parallel-size", type=int, default=1) temp_args, _ = temp_parser.parse_known_args() pp_size = temp_args.vllm_pipeline_parallel_size - vllm_tp_size = temp_args.rollout_num_gpus_per_engine // pp_size + pcp_size = temp_args.vllm_prefill_context_parallel_size + dp_size = temp_args.vllm_data_parallel_size + vllm_tp_size = temp_args.rollout_num_gpus_per_engine // (pp_size * pcp_size * dp_size) parser.set_defaults(vllm_tensor_parallel_size=vllm_tp_size) args, _ = parser.parse_known_args() - user_provided, raw_values = _detect_user_provided_dests(parser, sys.argv[1:]) - _sanitize_non_primitive_defaults(args, user_provided) - args._vllm_user_provided = user_provided - args._vllm_raw_values = raw_values return args - - -# Dests that are vime orchestration flags (not part of `vllm serve` CLI) — excluded -# from get_vllm_cli_action_table() so launch_server_process won't forward them. -_VIME_ORCHESTRATION_DESTS = frozenset( - { - "vllm_router_ip", - "vllm_router_port", - "router_request_timeout_secs", - "router_policy", - "vllm_server_concurrency", - "vllm_enable_deterministic_inference", - "vllm_weight_sync_packed", - "vllm_tool_call_parser", - "vllm_config", - "prefill_num_servers", - } -) - - -_VLLM_CLI_ACTION_TABLE_CACHE: dict[str, tuple[str, argparse.Action]] | None = None - - -def get_vllm_cli_action_table(): - """Build {vime_dest -> (primary_flag, action)} mapping for forwardable flags. - - Used by ``vllm_engine.launch_server_process`` to forward ``args.vllm_*`` values - that differ from vllm-side defaults to the ``vllm serve`` subprocess. - - Excludes vime orchestration dests and non-vllm-prefixed actions. Cached after - first build — rebuilding the parser is expensive. - """ - global _VLLM_CLI_ACTION_TABLE_CACHE - if _VLLM_CLI_ACTION_TABLE_CACHE is not None: - return _VLLM_CLI_ACTION_TABLE_CACHE - - parser = FlexibleArgumentParser(add_help=False) - add_vllm_arguments(parser) - - table: dict[str, tuple[str, argparse.Action]] = {} - for action in parser._actions: - if action.dest in _VIME_ORCHESTRATION_DESTS: - continue - if not action.dest.startswith("vllm_"): - continue - # Pick the first ``--vllm-xxx`` flag (skip ``--no-vllm-xxx`` companions). - primary_flag = None - for s in action.option_strings: - if s.startswith("--vllm-") and not s.startswith("--no-vllm-"): - primary_flag = "--" + s[len("--vllm-") :] - break - if primary_flag is None: - continue - table[action.dest] = (primary_flag, action) - _VLLM_CLI_ACTION_TABLE_CACHE = table - return table diff --git a/vime/backends/vllm_utils/deployment.py b/vime/backends/vllm_utils/deployment.py new file mode 100644 index 000000000..3fa1e4e87 --- /dev/null +++ b/vime/backends/vllm_utils/deployment.py @@ -0,0 +1,168 @@ +"""Launch and connect vLLM rollout deployments.""" + +import logging +import multiprocessing +import random +import time +from typing import Any + +from vime.backends.vllm_utils.disaggregation import collect_pd_urls, start_epd_server_groups, start_pd_server_groups +from vime.backends.vllm_utils.engine_group import RolloutServer, ServerGroupPlacement +from vime.backends.vllm_utils.external import start_external_rollout_servers +from vime.backends.vllm_utils.vllm_config import resolve_vllm_config +from vime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info + +logger = logging.getLogger(__name__) + + +def _start_router( + args, + *, + has_pd_disaggregation: bool = False, + force_new: bool = False, + bind: tuple[str, int] | None = None, + prefill_urls: list | None = None, + decode_urls: list | None = None, +) -> tuple[str, int, int | None]: + """Start the rollout HTTP gateway (vllm-router).""" + if bind is not None: + router_ip, router_port = bind + else: + if not force_new and args.vllm_router_ip is not None: + return args.vllm_router_ip, args.vllm_router_port, None + router_ip = _wrap_ipv6(get_host_info()[1]) + if force_new or args.vllm_router_port is None: + router_port = find_available_port(random.randint(3000, 4000)) + else: + router_port = args.vllm_router_port + + from vllm_router.router_args import RouterArgs + + from vime.utils.http_utils import run_router + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + router_args.host = router_ip + router_args.port = router_port + router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) + router_args.log_level = "warning" + router_args.request_timeout_secs = args.vllm_router_request_timeout_secs + + if has_pd_disaggregation: + router_args.vllm_pd_disaggregation = True + + if prefill_urls is not None: + router_args.prefill_urls = prefill_urls + router_args.decode_urls = decode_urls + + # Disable circuit breaker to prevent RDMA transfer timeouts from + # marking workers as dead. Timeouts are transient (PCIe contention under + # high load) and do not indicate a dead server. + router_args.disable_circuit_breaker = True + + logger.info(f"Launch router with args: {router_args}") + + process = multiprocessing.Process(target=run_router, args=(router_args,)) + process.daemon = True + process.start() + time.sleep(3) + assert process.is_alive() + logger.info(f"Router launched at {router_ip}:{router_port}, Prometheus port: {router_args.prometheus_port}") + return router_ip, router_port, router_args.prometheus_port + + +def _compute_rollout_offset(args) -> int: + """Offset (in PG bundle slots) where rollout GPUs start.""" + if args.debug_train_only or args.debug_rollout_only or args.colocate: + return 0 + offset = args.actor_num_nodes * args.actor_num_gpus_per_node + return offset + + +def _compute_megatron_num_gpus(args) -> int: + """Total number of megatron (actor + critic) GPU slots in the placement group.""" + if args.debug_rollout_only: + return 0 + num = args.actor_num_nodes * args.actor_num_gpus_per_node + return num + + +def start_rollout_servers(args, pg) -> tuple[dict[str, Any], list[Any]]: + """Start configured rollout servers without waiting for final engine initialization.""" + if args.rollout_external: + return start_external_rollout_servers(args, start_router=_start_router) + + config = resolve_vllm_config(args) + placement = ServerGroupPlacement( + args=args, + pg=pg, + rollout_pg_offset=_compute_rollout_offset(args), + megatron_num_gpus=_compute_megatron_num_gpus(args), + ) + + servers: dict[str, RolloutServer] = {} + encoder_metadata: dict[str, tuple[str, list[str]]] = {} + pending_init_handles: list[Any] = [] + + for model_idx, model_config in enumerate(config.models): + model_config.resolve(args) + has_pd = model_config.has_pd_disaggregation + + if has_pd: + router_ip = _wrap_ipv6(get_host_info()[1]) + router_port = find_available_port(random.randint(3000, 4000)) + prometheus_port = None + engine_router_ip = engine_router_port = None + else: + router_ip, router_port, prometheus_port = _start_router( + args, + force_new=(model_idx > 0), + ) + engine_router_ip, engine_router_port = router_ip, router_port + + if model_idx == 0: + args.vllm_router_ip = router_ip + args.vllm_router_port = router_port + + if model_config.has_encoder_disaggregation: + server_groups, init_handles, encoder_endpoints = start_epd_server_groups( + model_config, + placement, + engine_router_ip, + engine_router_port, + ) + encoder_metadata[model_config.name] = ( + server_groups[0].model_path, + encoder_endpoints, + ) + else: + server_groups, init_handles = start_pd_server_groups( + model_config, + placement, + engine_router_ip, + engine_router_port, + ) + + pending_init_handles.extend(init_handles) + + if has_pd: + prefill_urls, decode_urls = collect_pd_urls(server_groups) + _, _, prometheus_port = _start_router( + args, + has_pd_disaggregation=True, + bind=(router_ip, router_port), + prefill_urls=prefill_urls, + decode_urls=decode_urls, + ) + + servers[model_config.name] = RolloutServer( + server_groups=server_groups, + router_ip=router_ip, + router_port=router_port, + prometheus_port=prometheus_port, + model_name=model_config.name, + update_weights=model_config.update_weights, + ) + + args.vllm_model_routers = {name: (server.router_ip, server.router_port) for name, server in servers.items()} + args.vllm_model_encoder_endpoints = encoder_metadata + return servers, pending_init_handles diff --git a/vime/backends/vllm_utils/disaggregation.py b/vime/backends/vllm_utils/disaggregation.py new file mode 100644 index 000000000..2c6cfcc02 --- /dev/null +++ b/vime/backends/vllm_utils/disaggregation.py @@ -0,0 +1,104 @@ +"""PD and EPD-specific vLLM deployment sequencing.""" + +import logging +import uuid +from typing import Any + +import ray + +from vime.backends.vllm_utils.engine_group import ServerGroup, ServerGroupPlacement +from vime.backends.vllm_utils.vllm_config import ModelConfig + +logger = logging.getLogger(__name__) + + +def start_pd_server_groups( + model_config: ModelConfig, + placement: ServerGroupPlacement, + router_ip: str | None, + router_port: int | None, +) -> tuple[list[ServerGroup], list[Any]]: + """Start prefill/decode groups without waiting for final engine initialization.""" + server_groups = [] + init_handles = [] + for group_config in model_config.server_groups: + group = placement.create(group_config, router_ip, router_port) + init_handles.extend(placement.start(group)) + server_groups.append(group) + return server_groups, init_handles + + +def start_epd_server_groups( + model_config: ModelConfig, + placement: ServerGroupPlacement, + router_ip: str | None, + router_port: int | None, +) -> tuple[list[ServerGroup], list[Any], list[str]]: + """Start encoder groups first, then inject their endpoints into LLM groups.""" + server_groups = [] + transfer_overrides = { + "ec_transfer_config": { + "ec_connector_extra_config": { + "shared_storage_path": f"/dev/shm/vime-ec-{uuid.uuid4().hex}", + }, + }, + } + + encoder_endpoints: list[str] = [] + for group_config in model_config.server_groups: + if group_config.worker_type != "encoder": + continue + group = placement.create( + group_config, + router_ip, + router_port, + overrides_extra=transfer_overrides, + ) + handles = placement.start(group) + if handles: + ray.get(handles) + endpoints = ray.get([engine.get_url.remote() for engine in group.engines]) + encoder_endpoints.extend(endpoint for endpoint in endpoints if endpoint is not None) + server_groups.append(group) + + logger.info("EPD phase 1 done: collected %d encoder endpoints", len(encoder_endpoints)) + + init_handles = [] + for group_config in model_config.server_groups: + if group_config.worker_type == "encoder": + continue + overrides_extra = transfer_overrides if group_config.worker_type in ("regular", "prefill") else None + if overrides_extra is not None and encoder_endpoints: + overrides_extra = { + **transfer_overrides, + "language_only": True, + "encoder_urls": encoder_endpoints, + } + group = placement.create( + group_config, + router_ip, + router_port, + overrides_extra=overrides_extra, + ) + init_handles.extend(placement.start(group)) + server_groups.append(group) + + return server_groups, init_handles, encoder_endpoints + + +def collect_pd_urls(server_groups: list[ServerGroup]) -> tuple[list[tuple[str, None]], list[str]]: + """Collect static prefill/decode endpoints after engine initialization.""" + prefill_urls = [] + decode_urls = [] + for group in server_groups: + for engine in group.engines: + if engine is None: + continue + url = ray.get(engine.get_url.remote()) + if not url: + continue + if group.worker_type == "prefill": + prefill_urls.append((url, None)) + elif group.worker_type == "decode": + decode_urls.append(url) + return prefill_urls, decode_urls diff --git a/vime/backends/vllm_utils/engine_group.py b/vime/backends/vllm_utils/engine_group.py new file mode 100644 index 000000000..4142c7859 --- /dev/null +++ b/vime/backends/vllm_utils/engine_group.py @@ -0,0 +1,512 @@ +"""vLLM engine groups and their rollout-facing lifecycle.""" + +import dataclasses +import logging +import os +from typing import Any + +import ray +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from vime.backends.vllm_utils.vllm_config import ServerGroupConfig +from vime.backends.vllm_utils.vllm_engine import VLLMEngine, _resolve_parallel_sizes +from vime.platforms import current_platform +from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars + +GPU_MEMORY_TYPE_KV_CACHE = "kv_cache" +GPU_MEMORY_TYPE_WEIGHTS = "weights" +GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ServerGroup: + """A group of homogeneous vLLM engines with the same configuration. + + All engines in a group share the same tp_size / nodes_per_engine / pg. + A RolloutServer may contain multiple ServerGroups (e.g. prefill vs decode + in PD disaggregation). + """ + + args: Any + pg: Any # (placement_group, reordered_bundle_indices, reordered_gpu_ids) + all_engines: list + num_gpus_per_engine: int + num_new_engines: int + worker_type: str = "regular" # "regular", "prefill", "decode", or "placeholder" + rank_offset: int = 0 # cumulative engine count before this group + gpu_offset: int = 0 # cumulative GPU count before this group + vllm_overrides: dict = dataclasses.field(default_factory=dict) + needs_offload: bool = False # True when this group's GPUs overlap with megatron + model_path: str | None = None # checkpoint path for update_weights_from_disk + router_ip: str | None = None + router_port: int | None = None + + @property + def nodes_per_engine(self): + return max(1, self.num_gpus_per_engine // self.args.num_gpus_per_node) + + @property + def engines(self): + """Node-0 engines only (for multi-node serving).""" + return self.all_engines[:: self.nodes_per_engine] + + def parallel_config(self) -> dict[str, Any]: + """Return the VLLM parallel args that affect rank-local expert routing.""" + overrides = {key.replace("-", "_"): value for key, value in self.vllm_overrides.items()} + tp_size, pp_size, pcp_size, dp_size = _resolve_parallel_sizes( + self.args, + gpus_per_engine=self.num_gpus_per_engine, + overrides=overrides, + ) + enable_expert_parallel = bool( + overrides.get( + "enable_expert_parallel", + getattr(self.args, "vllm_enable_expert_parallel", False), + ) + ) + return { + "tp_size": tp_size, + "pp_size": pp_size, + "pcp_size": pcp_size, + "dp_size": dp_size, + "enable_expert_parallel": enable_expert_parallel, + "ep_size": tp_size * pcp_size * dp_size if enable_expert_parallel else 1, + } + + def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[list, dict[int, int]]: + """Create Ray actors, allocate ports, and fire ``engine.init()`` without waiting. + + Returns ``(init_handles, port_cursors)`` where *init_handles* is a list + of Ray ObjectRefs and *port_cursors* maps node index → next free port. + The caller should ``ray.get()`` on the handles to block until the + engines are healthy, and pass *port_cursors* to the next server group + so that different groups on the same node don't race for ports. + + Placeholder groups (worker_type="placeholder") skip engine creation entirely. + """ + if port_cursors is None: + port_cursors = {} + if self.args.debug_train_only or self.worker_type == "placeholder": + self.num_new_engines = 0 + return [], port_cursors + + num_gpus_per_engine_on_node = min(self.num_gpus_per_engine, self.args.num_gpus_per_node) + + pg, reordered_bundle_indices, reordered_gpu_ids = self.pg + num_engines = len(self.all_engines) + required_gpu_slots = self.gpu_offset + num_engines * num_gpus_per_engine_on_node + if num_engines and not ( + self.gpu_offset >= 0 and num_gpus_per_engine_on_node > 0 and required_gpu_slots <= len(reordered_gpu_ids) + ): + raise ValueError( + "Invalid rollout server group GPU placement: " + f"worker_type={self.worker_type}, " + f"gpu_offset={self.gpu_offset}, " + f"num_gpus_per_engine={self.num_gpus_per_engine}, " + f"num_gpus_per_engine_on_node={num_gpus_per_engine_on_node}, " + f"num_engines={num_engines}, " + f"required_gpu_slots={required_gpu_slots}, " + f"len(reordered_gpu_ids)={len(reordered_gpu_ids)}, " + f"rollout_num_gpus={self.args.rollout_num_gpus}, " + f"rollout_num_gpus_per_engine={self.args.rollout_num_gpus_per_engine}. " + "Please align --rollout-num-gpus, --rollout-num-gpus-per-engine, " + "and --vllm-config server_groups." + ) + + RolloutRayActor = ray.remote(VLLMEngine) + + rollout_engines = [] + platform = current_platform() + for i in range(len(self.all_engines)): + if self.all_engines[i] is not None: + continue + + global_rank = self.rank_offset + i + num_gpus = 0.2 + num_cpus = num_gpus + + # Get the base GPU ID from placement group using gpu_offset. + gpu_index = self.gpu_offset + i * num_gpus_per_engine_on_node + base_gpu_id = int(reordered_gpu_ids[gpu_index]) + + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[gpu_index], + ) + + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} + # vime-patch: expandable_segments breaks vLLM custom all-reduce CUDA + # IPC. Strip only that key, keeping any other allocator settings. + _alloc = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") + env_vars["PYTORCH_CUDA_ALLOC_CONF"] = ",".join( + kv for kv in _alloc.split(",") if kv and not kv.strip().startswith("expandable_segments") + ) + if platform.is_npu: + env_vars = platform.ray.rollout_runtime_env(self.args, env_vars) + resource_options = {"num_gpus": num_gpus} + if platform.is_npu: + resource_options = {"num_gpus": 0, **platform.ray.actor_options(num_gpus)} + rollout_engine = RolloutRayActor.options( + num_cpus=num_cpus, + scheduling_strategy=scheduling_strategy, + runtime_env={ + "env_vars": add_default_ray_env_vars(env_vars), + }, + **resource_options, + ).remote( + self.args, + rank=global_rank, + worker_type=self.worker_type, + base_gpu_id=base_gpu_id, + vllm_overrides=self.vllm_overrides, + num_gpus_per_engine=self.num_gpus_per_engine, + ) + + rollout_engines.append((global_rank, rollout_engine)) + self.all_engines[i] = rollout_engine + + self.num_new_engines = len(rollout_engines) + + if self.num_new_engines == 0: + return [], port_cursors + + # Compute base_port from the maximum cursor across all nodes that + # this group's engines may land on (conservative: just use global max). + base_port = max(port_cursors.values()) if port_cursors else 15000 + addr_and_ports, port_cursors = _allocate_rollout_engine_addr_and_ports_normal( + args=self.args, + rollout_engines=rollout_engines, + worker_type=self.worker_type, + num_gpus_per_engine=self.num_gpus_per_engine, + rank_offset=self.rank_offset, + base_port=base_port, + ) + + init_handles = [ + engine.init.remote( + **(addr_and_ports[rank]), + router_ip=self.router_ip, + router_port=self.router_port, + ) + for rank, engine in rollout_engines + ] + return init_handles, port_cursors + + def offload(self): + """Fire release_memory_occupation on all engines (non-blocking). + + Returns a list of Ray ObjectRefs. Skipped for groups that do not + overlap with megatron GPUs (``needs_offload=False``). + """ + if not self.needs_offload: + return [] + return [engine.release_memory_occupation.remote() for engine in self.engines if engine is not None] + + def onload(self, tags: list[str] | None = None): + """Fire resume_memory_occupation on all engines (non-blocking). + + Returns a list of Ray ObjectRefs. Skipped for groups that do not + overlap with megatron GPUs (``needs_offload=False``). + """ + if not self.needs_offload: + return [] + return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.engines if engine is not None] + + +@dataclasses.dataclass +class RolloutServer: + """A model served behind a shared router, with one or more server groups. + + Each RolloutServer represents one model deployed behind a single router. + A server may contain multiple ServerGroups with different + ``num_gpus_per_engine`` (e.g. prefill TP=2, decode TP=4). + """ + + server_groups: list[ServerGroup] + router_ip: str | None = None + router_port: int | None = None + prometheus_port: int | None = None + model_name: str = "default" + update_weights: bool = True + + @property + def engines(self): + """All node-0 engines across all groups (placeholder groups contribute nothing).""" + return [e for g in self.server_groups for e in g.engines] + + @property + def all_engines(self): + """All engines (including non-node-0) across all groups.""" + return [e for g in self.server_groups for e in g.all_engines] + + @property + def num_new_engines(self): + return sum(g.num_new_engines for g in self.server_groups) + + @num_new_engines.setter + def num_new_engines(self, value): + for g in self.server_groups: + g.num_new_engines = value + + @property + def engine_gpu_counts(self) -> list[int]: + """Per-engine GPU count for all node-0 engines, parallel to ``engines``.""" + return [g.num_gpus_per_engine for g in self.server_groups for _ in g.engines] + + @property + def engine_gpu_offsets(self) -> list[int]: + """Per-engine GPU offset for all node-0 engines, parallel to ``engines``. + + Accounts for placeholder groups that occupy GPU slots without creating engines. + """ + offsets = [] + for g in self.server_groups: + for j in range(len(g.engines)): + offsets.append(g.gpu_offset + j * g.num_gpus_per_engine) + return offsets + + @property + def engine_parallel_configs(self) -> list[dict[str, Any]]: + """Per-engine VLLM parallel config, parallel to ``engines``.""" + return [g.parallel_config() for g in self.server_groups for _ in g.engines] + + @property + def nodes_per_engine(self): + """Nodes per engine. Only valid when all active groups share the same value.""" + values = {g.nodes_per_engine for g in self.server_groups if g.worker_type != "placeholder"} + if len(values) != 1: + raise ValueError(f"Heterogeneous nodes_per_engine across groups: {values}") + return values.pop() + + def recover(self): + """Recover dead engines across all active groups, overlapping init.""" + # Record dead indices per group before starting. + dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in self.server_groups] + + # Start all groups concurrently. + all_handles = [] + port_cursors: dict[int, int] = {} + for g in self.server_groups: + handles, port_cursors = g.start_engines(port_cursors) + all_handles.extend(handles) + if all_handles: + ray.get(all_handles) + + # Post-recovery: offload then onload weights for newly created engines. + release_handles = [] + updatable_new_engines = [] + non_updatable_groups_engines: list[tuple[str, list]] = [] + for g, dead_indices in zip(self.server_groups, dead_per_group, strict=True): + logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})") + assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length" + if g.needs_offload and dead_indices: + new_engines = [g.all_engines[i] for i in dead_indices] + release_handles.extend(engine.release_memory_occupation.remote() for engine in new_engines) + if self.update_weights: + updatable_new_engines.extend(new_engines) + elif g.model_path: + non_updatable_groups_engines.append((g.model_path, new_engines)) + + if release_handles: + ray.get(release_handles) + # Resume GPU memory for all engines that need offload. + all_resume_engines = updatable_new_engines[:] + for _model_path, engines in non_updatable_groups_engines: + all_resume_engines.extend(engines) + if all_resume_engines: + ray.get( + [ + engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + for engine in all_resume_engines + ] + ) + + def offload(self): + """Release memory occupation across all groups (concurrent).""" + handles = [] + for g in self.server_groups: + handles.extend(g.offload()) + return ray.get(handles) if handles else [] + + def onload(self, tags: list[str] | None = None): + """Resume memory occupation across all groups (concurrent).""" + handles = [] + for g in self.server_groups: + handles.extend(g.onload(tags)) + return ray.get(handles) if handles else [] + + def onload_weights(self): + """Restore weights for offloaded groups. + + All groups resume from CPU cache via ``resume_memory_occupation``. + For updatable servers, weights will be overwritten by + ``update_weights`` shortly after. For non-updatable servers the + CPU backup already contains the correct (unchanged) weights. + """ + handles = [] + for g in self.server_groups: + if not g.needs_offload: + continue + handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])) + return ray.get(handles) if handles else [] + + def onload_kv(self): + """Resume KV cache and CUDA graphs for offloaded groups.""" + handles = [] + for g in self.server_groups: + handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])) + return ray.get(handles) if handles else [] + + +@dataclasses.dataclass +class ServerGroupPlacement: + """Track rank and GPU offsets while materializing configured server groups.""" + + args: Any + pg: Any + rollout_pg_offset: int + megatron_num_gpus: int + engine_offset: int = 0 + gpu_offset: int = 0 + port_cursors: dict[int, int] = dataclasses.field(default_factory=dict) + + def create( + self, + group_config: ServerGroupConfig, + router_ip: str | None, + router_port: int | None, + *, + overrides_extra: dict | None = None, + ) -> ServerGroup: + gpus_per_engine = group_config.num_gpus_per_engine + assert gpus_per_engine is not None, "ModelConfig.resolve() must set num_gpus_per_engine before deployment" + num_gpus_per_engine_on_node = min(gpus_per_engine, self.args.num_gpus_per_node) + num_engines = group_config.num_gpus // num_gpus_per_engine_on_node + + group_abs_start = self.rollout_pg_offset + self.gpu_offset + needs_offload = self.args.offload_rollout and group_abs_start < self.megatron_num_gpus + overrides = dict(group_config.overrides) + if overrides_extra: + for key, value in overrides_extra.items(): + overrides.setdefault(key, value) + if self.args.offload_rollout and not needs_offload: + overrides.setdefault("enable_memory_saver", False) + logger.info( + f"Engine group '{group_config.worker_type}' gpu_offset={self.gpu_offset} " + f"(abs={group_abs_start}): needs_offload={needs_offload}" + ) + + group = ServerGroup( + args=self.args, + pg=self.pg, + all_engines=[None] * num_engines if group_config.worker_type != "placeholder" else [], + num_gpus_per_engine=gpus_per_engine, + num_new_engines=0, + worker_type=group_config.worker_type, + rank_offset=self.engine_offset, + gpu_offset=self.gpu_offset, + vllm_overrides=overrides, + needs_offload=needs_offload, + model_path=overrides.get("model_path", self.args.hf_checkpoint), + router_ip=router_ip, + router_port=router_port, + ) + self.engine_offset += num_engines + self.gpu_offset += group_config.num_gpus + return group + + def start(self, group: ServerGroup) -> list: + handles, self.port_cursors = group.start_engines(self.port_cursors) + return handles + + +def _allocate_rollout_engine_addr_and_ports_normal( + *, + args, + rollout_engines, + worker_type="regular", + num_gpus_per_engine=None, + rank_offset=0, + base_port=15000, +): + # get ports + # there are 4 ports we need to allocate + # 1. server port + # 2. nccl port + # 3. dist_init_addr port + # 4. other ports for dp_attention, which is of size 4 + dp_size + _gpus_per_engine = num_gpus_per_engine or args.rollout_num_gpus_per_engine + num_engines_per_node = max(1, args.num_gpus_per_node // _gpus_per_engine) + addr_and_ports: dict[int, dict] = {} + + # Track per-node port cursors so that different server groups (called + # sequentially) never race for the same ports on a given node. + node_port_cursor: dict[int, int] = {} + + visited_nodes = set() + for rank, engine in rollout_engines: + local_rank = rank - rank_offset + node_index = local_rank // num_engines_per_node + if node_index in visited_nodes: + continue + visited_nodes.add(node_index) + # TODO: currently when restarting engines, we will set port for all engines on this node starting with this rank. + # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node. + num_engines_on_this_node = num_engines_per_node - (local_rank % num_engines_per_node) + + def get_addr_and_ports(engine, node_idx): + # use small ports to prevent ephemeral port between 32768 and 65536. + # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition + start_port = node_port_cursor.get(node_idx, base_port) + + def port(consecutive=1): + nonlocal start_port + _, port = ray.get( + engine._get_current_node_ip_and_free_port.remote( + start_port=start_port, + consecutive=consecutive, + ) + ) + start_port = port + consecutive + node_port_cursor[node_idx] = start_port + return port + + def addr(): + addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) + return addr + + return addr, port + + get_addr, get_port = get_addr_and_ports(engine, node_index) + + for i in range(num_engines_on_this_node): + current_rank = rank + i + addr_and_ports.setdefault(current_rank, {}) + addr_and_ports[current_rank]["host"] = get_addr() + addr_and_ports[current_rank]["port"] = get_port() + addr_and_ports[current_rank]["nccl_port"] = get_port() + + if worker_type in ("prefill", "decode"): + addr_and_ports[current_rank]["disaggregation_bootstrap_port"] = get_port() + + if _gpus_per_engine > args.num_gpus_per_node: + num_node_per_engine = _gpus_per_engine // args.num_gpus_per_node + if local_rank % num_node_per_engine == 0: + # this is the first node in the engine, we need to allocate the dist_init_addr port + dist_init_addr = f"{get_addr()}:{get_port(30 + args.vllm_dp_size)}" + for i in range(num_node_per_engine): + addr_and_ports.setdefault(rank + i, {}) + addr_and_ports[rank + i]["dist_init_addr"] = dist_init_addr + else: + for i in range(num_engines_on_this_node): + addr_and_ports[rank + i]["dist_init_addr"] = f"{get_addr()}:{get_port(30 + args.vllm_dp_size)}" + + for i, _ in rollout_engines: + for key in ["port", "nccl_port", "dist_init_addr"]: + assert key in addr_and_ports[i], f"Engine {i} {key} is not set." + logger.info(f"Ports for engine {i}: {addr_and_ports[i]}") + + return addr_and_ports, node_port_cursor diff --git a/vime/backends/vllm_utils/external.py b/vime/backends/vllm_utils/external.py new file mode 100644 index 000000000..bb78dd034 --- /dev/null +++ b/vime/backends/vllm_utils/external.py @@ -0,0 +1,320 @@ +"""Helpers for pre-launched external vLLM engines.""" + +from __future__ import annotations + +import dataclasses +import logging +from urllib.parse import urlparse + +import requests + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class ExternalEngineInfo: + url: str + host: str + port: int + worker_type: str + num_gpus: int + disaggregation_bootstrap_port: int | None = None + server_info: dict = dataclasses.field(default_factory=dict) + + @property + def is_pd_worker(self) -> bool: + return self.worker_type in ("prefill", "decode") + + @property + def parallel_config(self) -> dict[str, int | bool]: + pp_size = int(self.server_info.get("pp_size") or self.server_info.get("pipeline_parallel_size") or 1) + pcp_size = int(self.server_info.get("pcp_size") or self.server_info.get("prefill_context_parallel_size") or 1) + dp_size = int(self.server_info.get("dp_size") or self.server_info.get("data_parallel_size") or 1) + tp_size = int( + self.server_info.get("tp_size") + or self.server_info.get("tensor_parallel_size") + or self.num_gpus // (pp_size * pcp_size * dp_size) + ) + enable_expert_parallel = bool(self.server_info.get("enable_expert_parallel", False)) + return { + "tp_size": tp_size, + "pp_size": pp_size, + "pcp_size": pcp_size, + "dp_size": dp_size, + "enable_expert_parallel": enable_expert_parallel, + "ep_size": tp_size * pcp_size * dp_size if enable_expert_parallel else 1, + } + + def to_dict(self) -> dict: + return dataclasses.asdict(self) + + +def normalize_external_engine_addr(addr: str) -> str: + """Normalize ``host:port`` or ``http://host:port`` to an HTTP base URL.""" + if "://" not in addr: + addr = f"http://{addr}" + addr = addr.rstrip("/") + parsed = urlparse(addr) + if parsed.scheme != "http" or parsed.hostname is None or parsed.port is None: + raise ValueError( + f"Invalid external vLLM engine address {addr!r}. " + "Use host:port or http://host:port (IPv6 must be bracketed)." + ) + return addr + + +def external_engine_init_kwargs(info: ExternalEngineInfo) -> dict: + init_kwargs = { + "dist_init_addr": f"{info.host}:{info.port}", + "nccl_port": None, + "host": info.host, + "port": info.port, + } + if info.worker_type == "prefill": + init_kwargs["disaggregation_bootstrap_port"] = info.disaggregation_bootstrap_port + return init_kwargs + + +def get_server_info(url: str, timeout: float = 30.0) -> dict: + errors = [] + for endpoint in ("/server_info?config_format=json", "/server_info", "/get_server_info"): + try: + response = requests.get(f"{url}{endpoint}", timeout=timeout) + response.raise_for_status() + return _normalize_server_info(response.json()) + except Exception as exc: + errors.append(f"{endpoint}: {exc}") + raise RuntimeError(f"Failed to fetch vLLM server info from {url}: {'; '.join(errors)}") + + +def _normalize_server_info(server_info: dict) -> dict: + vllm_config = server_info.get("vllm_config") + if not isinstance(vllm_config, dict): + return server_info + + normalized = dict(server_info) + for section in vllm_config.values(): + if not isinstance(section, dict): + continue + for key, value in section.items(): + normalized.setdefault(key, value) + + def find_config_value(config, name): + if not isinstance(config, dict): + return None + if name in config: + return config[name] + for value in config.values(): + found = find_config_value(value, name) + if found is not None: + return found + return None + + kv_transfer_config = find_config_value(vllm_config, "kv_transfer_config") + if kv_transfer_config is not None: + normalized["kv_transfer_config"] = kv_transfer_config + weight_transfer_config = find_config_value(vllm_config, "weight_transfer_config") + if weight_transfer_config is not None: + normalized["weight_transfer_config"] = weight_transfer_config + if isinstance(kv_transfer_config, dict): + role = kv_transfer_config.get("kv_role") + if role == "kv_producer": + normalized["disaggregation_mode"] = "prefill" + elif role == "kv_consumer": + normalized["disaggregation_mode"] = "decode" + + vllm_env = server_info.get("vllm_env") + if isinstance(vllm_env, dict): + bootstrap_port = vllm_env.get("VLLM_NIXL_SIDE_CHANNEL_PORT") + if bootstrap_port is not None: + normalized["disaggregation_bootstrap_port"] = bootstrap_port + + return normalized + + +def _infer_worker_type(server_info: dict) -> str: + if server_info.get("encoder_only"): + return "encoder" + kv_transfer_config = server_info.get("kv_transfer_config") + if isinstance(kv_transfer_config, dict): + role = kv_transfer_config.get("kv_role") + if role == "kv_producer": + return "prefill" + if role == "kv_consumer": + return "decode" + mode = server_info.get("disaggregation_mode") + if mode in ("prefill", "decode"): + return mode + return "regular" + + +def discover_external_engines(addrs: list[str], timeout: float = 30.0) -> list[ExternalEngineInfo]: + infos = [] + for addr in addrs: + url = normalize_external_engine_addr(addr) + parsed = urlparse(url) + assert parsed.hostname is not None and parsed.port is not None + server_info = get_server_info(url, timeout=timeout) + + pp_size = int(server_info.get("pp_size") or server_info.get("pipeline_parallel_size") or 1) + pcp_size = int(server_info.get("pcp_size") or server_info.get("prefill_context_parallel_size") or 1) + dp_size = int(server_info.get("dp_size") or server_info.get("data_parallel_size") or 1) + tp_size = int(server_info.get("tp_size") or server_info.get("tensor_parallel_size") or 1) + num_gpus = int( + server_info.get("num_gpus") + or server_info.get("num_gpus_per_engine") + or tp_size * pp_size * pcp_size * dp_size + ) + bootstrap_port = server_info.get("disaggregation_bootstrap_port") + bootstrap_port = int(bootstrap_port) if bootstrap_port is not None else None + + infos.append( + ExternalEngineInfo( + url=url, + host=parsed.hostname, + port=parsed.port, + worker_type=_infer_worker_type(server_info), + num_gpus=num_gpus, + disaggregation_bootstrap_port=bootstrap_port, + server_info=server_info, + ) + ) + return infos + + +def apply_external_engine_info_to_args(args, logger=None) -> None: + """Detect external engines and store the derived topology on ``args``.""" + addrs = args.rollout_external_engine_addrs + if not addrs: + raise ValueError("apply_external_engine_info_to_args requires --rollout-external-engine-addrs.") + + infos = discover_external_engines(addrs) + if not infos: + raise ValueError("--rollout-external-engine-addrs did not contain any engines.") + + args.rollout_external_engine_infos = [info.to_dict() for info in infos] + args.rollout_num_engines = len(infos) + args.rollout_num_gpus = sum(info.num_gpus for info in infos) + + if logger is not None: + summary = [ + { + "url": info.url, + "worker_type": info.worker_type, + "num_gpus": info.num_gpus, + "disaggregation_bootstrap_port": info.disaggregation_bootstrap_port, + } + for info in infos + ] + logger.info(f"Detected external vLLM engines: {summary}") + + +@dataclasses.dataclass +class ExternalRolloutServer: + """Rollout server backed by pre-launched external vLLM engines.""" + + engines: list + engine_gpu_counts: list[int] + engine_gpu_offsets: list[int] + engine_parallel_configs: list[dict[str, int]] + router_ip: str | None = None + router_port: int | None = None + model_name: str = "default" + update_weights: bool = True + num_new_engines: int = 0 + server_groups: list = dataclasses.field(default_factory=list) + + @property + def all_engines(self): + return self.engines + + def recover(self): + logger.warning("Fault tolerance is not supported for external rollout engines; skip recover.") + + def offload(self): + return [] + + def onload(self, tags: list[str] | None = None): + return [] + + def onload_weights(self): + return [] + + def onload_kv(self): + return [] + + +def external_engine_infos_from_args(args) -> list[ExternalEngineInfo]: + raw_infos = getattr(args, "rollout_external_engine_infos", None) + if raw_infos is None: + raise RuntimeError( + "External rollout engine info is missing. " + "apply_external_engine_info_to_args must run before starting external rollout servers." + ) + return [ExternalEngineInfo(**info) if isinstance(info, dict) else info for info in raw_infos] + + +def start_external_rollout_servers(args, *, start_router) -> tuple[dict[str, ExternalRolloutServer], list]: + import ray + + from vime.backends.vllm_utils.vllm_engine import VLLMEngine + from vime.ray.utils import add_default_ray_env_vars + + infos = external_engine_infos_from_args(args) + has_pd_disaggregation = any(info.is_pd_worker for info in infos) + prefill_urls = [(info.url, info.disaggregation_bootstrap_port) for info in infos if info.worker_type == "prefill"] + decode_urls = [info.url for info in infos if info.worker_type == "decode"] + router_ip, router_port, _ = start_router( + args, + has_pd_disaggregation=has_pd_disaggregation, + prefill_urls=prefill_urls if has_pd_disaggregation else None, + decode_urls=decode_urls if has_pd_disaggregation else None, + ) + args.vllm_router_ip = router_ip + args.vllm_router_port = router_port + + engines = [] + engine_gpu_counts = [] + engine_gpu_offsets = [] + init_handles = [] + RolloutRayActor = ray.remote(VLLMEngine) + gpu_offset = 0 + for rank, info in enumerate(infos): + rollout_engine = RolloutRayActor.options( + num_cpus=0.2, + num_gpus=0, + runtime_env={"env_vars": add_default_ray_env_vars()}, + ).remote( + args=args, + rank=rank, + worker_type=info.worker_type, + base_gpu_id=0, + num_gpus_per_engine=info.num_gpus, + ) + engines.append(rollout_engine) + engine_gpu_counts.append(info.num_gpus) + engine_gpu_offsets.append(gpu_offset) + gpu_offset += info.num_gpus + init_handles.append( + rollout_engine.init.remote( + **external_engine_init_kwargs(info), + router_ip=None if has_pd_disaggregation else router_ip, + router_port=None if has_pd_disaggregation else router_port, + ) + ) + + args.vllm_model_routers = {"default": (router_ip, router_port)} + servers = { + "default": ExternalRolloutServer( + engines=engines, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + engine_parallel_configs=[info.parallel_config for info in infos], + router_ip=router_ip, + router_port=router_port, + model_name="default", + update_weights=True, + num_new_engines=len(engines), + ) + } + return servers, init_handles diff --git a/vime/backends/vllm_utils/server_control.py b/vime/backends/vllm_utils/server_control.py new file mode 100644 index 000000000..99b45e8cd --- /dev/null +++ b/vime/backends/vllm_utils/server_control.py @@ -0,0 +1,24 @@ +"""Control-plane helper for aborting in-flight requests on vLLM workers.""" + +import asyncio +import logging + +from vime.utils.http_utils import post + +logger = logging.getLogger(__name__) + + +async def abort_inflight_requests(urls: list[str]) -> None: + """Abort all in-flight requests on each worker (one best-effort sweep). + + Posts to ``/abort_requests`` with an empty body; failures are logged, not + raised. Idempotent, so the caller may re-issue it to converge. + """ + + async def _abort_one(url: str) -> None: + try: + await post(f"{url.rstrip('/')}/abort_requests", {}, max_retries=3) + except Exception as e: + logger.warning(f"Failed to abort requests on {url}: {e}") + + await asyncio.gather(*(_abort_one(url) for url in urls)) diff --git a/vime/backends/vllm_utils/vllm_config.py b/vime/backends/vllm_utils/vllm_config.py index 4dd02bceb..2fcca37ad 100644 --- a/vime/backends/vllm_utils/vllm_config.py +++ b/vime/backends/vllm_utils/vllm_config.py @@ -23,9 +23,15 @@ class ServerGroupConfig: num_gpus: Total number of GPUs for this group. num_gpus_per_engine: GPUs per engine for this group. Overrides the model-level or global ``--rollout-num-gpus-per-engine``. - overrides: Optional dict of vLLM ``ServerArgs`` field overrides. - These are applied on top of the base CLI ``--vllm-*`` - arguments in ``_compute_server_args``. + overrides: Optional dict of vLLM engine-arg field overrides, applied + on top of the base CLI ``--vllm-*`` arguments in + ``_compute_server_args`` (highest priority). Keys are vLLM + ``AsyncEngineArgs`` / ``FrontendArgs`` field names in + underscore style (e.g. ``gpu_memory_utilization``); the + exact accepted set is ``_vllm_server_field_names()`` in + ``vllm_engine``. The accepted set combines engine config + (``AsyncEngineArgs``) and OpenAI-frontend config + (``FrontendArgs``). """ worker_type: str @@ -140,8 +146,10 @@ class VllmConfig: num_gpus: 4 Each model gets its own router. ``placeholder`` groups reserve GPU - slots without creating engines. ``overrides`` are ``ServerArgs`` - field names applied on top of the base ``--vllm-*`` CLI args. + slots without creating engines. ``overrides`` are vLLM + ``AsyncEngineArgs`` / ``FrontendArgs`` field names (see + ``ServerGroupConfig.overrides`` and ``vllm_engine._vllm_server_field_names``) + applied on top of the base ``--vllm-*`` CLI args. Set ``update_weights: false`` for frozen models (reference, reward, etc.) that should not receive weight updates from training. @@ -205,3 +213,28 @@ def has_pd_disaggregation(self) -> bool: @property def total_num_gpus(self) -> int: return sum(m.total_num_gpus for m in self.models) + + +def resolve_vllm_config(args) -> VllmConfig: + """Resolve the configured, legacy PD, or default vLLM deployment.""" + if getattr(args, "vllm_config", None) is not None: + config = VllmConfig.from_yaml(args.vllm_config) + expected = args.rollout_num_gpus + actual = config.total_num_gpus + assert actual == expected, f"vllm_config total GPUs ({actual}) != rollout_num_gpus ({expected})" + return config + + if args.rollout_num_gpus == 0: + return VllmConfig(models=[ModelConfig(name="default", server_groups=[])]) + + if args.prefill_num_servers is not None: + return VllmConfig.from_prefill_num_servers(args) + + return VllmConfig( + models=[ + ModelConfig( + name="default", + server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], + ) + ] + ) diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 69de3a1b7..eaa4efddc 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -1,72 +1,30 @@ -"""Ray actor and launch helpers for vLLM OpenAI HTTP rollout. - -Per-Ray-actor ``server_args`` dict is built via :func:`_compute_server_args`, -then :func:`build_vllm_cmd_and_env` turns it into ``vllm serve`` CLI + subprocess env. -:class:`VLLMEngine` manages the runtime HTTP control plane. -User-facing vLLM knobs remain on ``train.py`` as ``--vllm-*`` (see ``arguments.py``). -""" - -from __future__ import annotations - +import argparse import base64 import dataclasses -import ipaddress import logging import multiprocessing import os -import pickle import time -from argparse import BooleanOptionalAction from typing import Any from urllib.parse import quote +import cloudpickle import requests +from vllm.utils.system_utils import kill_process_tree -from vime.backends.vllm_utils.arguments import SKIPPED_DESTS, get_vllm_cli_action_table +from vime.backends.vllm_utils.external import get_server_info +from vime.platforms import current_platform from vime.ray.ray_actor import RayActor -from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath -from vime.utils.http_utils import get_host_info +from vime.utils.http_utils import _wrap_ipv6, get_host_info logger = logging.getLogger(__name__) -_spawn_ctx = multiprocessing.get_context("spawn") - -# Fields checked against external ``GET /server_info``. -EXTERNAL_ENGINE_CHECK_FIELDS = ("tp_size", "pp_size", "dp_size", "nnodes") - -_REDACTED_FLAGS = frozenset({"--hf-token"}) - -# vLLM sleep/wake only supports these tags (``cuda_graph`` is not supported). _VLLM_WAKE_TAGS = frozenset({"weights", "kv_cache"}) - -_PRIMITIVE_TYPES = (str, int, float, bool) - - -def _format_v6_uri(addr: str | None) -> str | None: - if not addr or addr.startswith("["): - return addr - try: - if ipaddress.ip_address(addr).version == 6: - return f"[{addr}]" - except ValueError: - pass - return addr - - -def _response_json(response: requests.Response) -> dict: - try: - response.raise_for_status() - except requests.exceptions.HTTPError as e: - e.add_note(f"{response.text=}") - raise - # vLLM sleep/wake endpoints may return 200 with an empty body. - if not response.content or not response.content.strip(): - return {"ok": True} - return response.json() +_LEGACY_VLLM_PARALLEL_FIELDS = frozenset({"tp_size", "pp_size", "pcp_size", "dp_size"}) +_INVALID_VLLM_PARALLEL_FIELDS = frozenset({"ep_size", "expert_parallel_size", "moe_dp_size", "moe_data_parallel_size"}) def get_base_gpu_id(args, rank): - """First local GPU index on this node for rollout engine *rank* (colocate vs actor[/critic]-offset layout).""" num_gpus = min(args.num_gpus_per_node, args.rollout_num_gpus_per_engine) if args.colocate: start_index = (rank * num_gpus) % args.num_gpus_per_node @@ -79,297 +37,44 @@ def get_base_gpu_id(args, rank): return start_index -def _to_local_gpu_id(physical_gpu_id: int) -> int: - if is_npu(): - cvd = os.environ.get("ASCEND_RT_VISIBLE_DEVICES") - else: - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if not cvd: - return physical_gpu_id - visible = [int(x) for x in cvd.split(",") if x.strip() != ""] - if physical_gpu_id in visible: - return visible.index(physical_gpu_id) - if 0 <= physical_gpu_id < len(visible): - return physical_gpu_id - raise RuntimeError( - f"GPU id {physical_gpu_id} is not valid under CUDA_VISIBLE_DEVICES={cvd}. " - f"Expected one of {visible} (physical) or 0..{len(visible)-1} (local)." - ) - - -@dataclasses.dataclass(frozen=True) -class VllmEngineTopology: - """Per-Ray-actor placement for one slice of a logical rollout engine.""" - - nnodes: int - node_rank: int - local_num_gpus: int - tensor_parallel_size: int - pipeline_parallel_size: int - - @property - def headless(self) -> bool: - return self.node_rank != 0 - - @property - def multi_node(self) -> bool: - return self.nnodes > 1 +def launch_server_process(server_args_dict: dict) -> multiprocessing.Process: + env = _build_subprocess_env(server_args_dict) + kwargs = {k: v for k, v in server_args_dict.items() if not k.startswith("_")} + host = _wrap_ipv6(kwargs.get("host") or "127.0.0.1") + kwargs["host"] = host.strip("[]") + logger.info("Launching vLLM server: %s", kwargs) + multiprocessing.set_start_method("spawn", force=True) + p = multiprocessing.Process(target=_run_vllm_server, args=(kwargs, env)) + p.start() -def _get_vllm_pp_size(args) -> int: - return int(getattr(args, "vllm_pipeline_parallel_size", 1) or 1) - - -def _get_vllm_dp_size(args) -> int: - return int(getattr(args, "vllm_dp_size", None) or getattr(args, "vllm_data_parallel_size", 1) or 1) - - -def _resolve_vllm_parallel_sizes(args, *, gpus_per_engine: int) -> tuple[int, int]: - # Derive TP per-engine from THIS engine's GPU count (matches upstream slime's - # sglang_engine: tp = _gpus_per_engine // pp). Deliberately does NOT consult a global - # ``args.vllm_tp_size``: validate_args used to set that from the *global* - # rollout_num_gpus_per_engine, which shadowed this per-engine value and made a - # heterogeneous per-group engine (e.g. a tp=2 group) launch with the global TP — - # desyncing the weight-transfer rendezvous (the 300s "3/4 clients joined" hang). - pp = _get_vllm_pp_size(args) - dp = _get_vllm_dp_size(args) - if gpus_per_engine % (pp * dp) != 0: - raise ValueError( - f"num_gpus_per_engine ({gpus_per_engine}) must be divisible by " - f"vllm_pipeline_parallel_size * vllm_data_parallel_size ({pp} * {dp} = {pp * dp})" - ) - tp = gpus_per_engine // (pp * dp) - return tp, pp - + if server_args_dict.get("node_rank", 0) != 0: + return p -def compute_vllm_engine_topology( - args, - global_rank: int, - *, - num_gpus_per_engine: int | None = None, -) -> VllmEngineTopology: - """Compute nnodes / node_rank / local GPU slice for one Ray actor.""" - gpus_per_engine = num_gpus_per_engine if num_gpus_per_engine is not None else args.rollout_num_gpus_per_engine - nnodes = max(1, gpus_per_engine // args.num_gpus_per_node) - node_rank = global_rank % nnodes - if nnodes == 1: - local_num_gpus = min(args.num_gpus_per_node, gpus_per_engine) - else: - if gpus_per_engine % nnodes != 0: - raise ValueError( - f"rollout_num_gpus_per_engine ({gpus_per_engine}) must be divisible by the number of " - f"nodes per engine ({nnodes})" - ) - local_num_gpus = gpus_per_engine // nnodes - tp, pp = _resolve_vllm_parallel_sizes(args, gpus_per_engine=gpus_per_engine) - return VllmEngineTopology( - nnodes=nnodes, - node_rank=node_rank, - local_num_gpus=local_num_gpus, - tensor_parallel_size=tp, - pipeline_parallel_size=pp, + _wait_server_healthy( + base_url=f"http://{host}:{server_args_dict['port']}", + is_process_alive=lambda: p.is_alive(), ) - -def parse_dist_init_addr(dist_init_addr: str) -> tuple[str, int]: - """Split ``host:port`` (IPv6-safe) into master host and port.""" - ip_part, port_part = dist_init_addr.rsplit(":", 1) - host = _format_v6_uri(ip_part) or ip_part - return host.strip("[]"), int(port_part) - - -def append_vllm_distributed_launch_flags( - cmd: list[str], - topology: VllmEngineTopology, - master: tuple[str, int], - args, -) -> None: - """Append vLLM multi-node flags when ``topology.multi_node`` (no-op for single-node).""" - if not topology.multi_node: - return - master_host, master_port = master - cmd += [ - "--nnodes", - str(topology.nnodes), - "--node-rank", - str(topology.node_rank), - "--master-addr", - master_host, - "--master-port", - str(master_port), - ] - if not _user_overrode(args, "vllm_data_parallel_backend"): - cmd += ["--data-parallel-backend", "mp"] - if not _user_overrode(args, "vllm_distributed_executor_backend"): - cmd += ["--distributed-executor-backend", "mp"] - if topology.headless: - cmd.append("--headless") - - -def _user_overrode(args, dest: str) -> bool: - user_provided: set[str] = getattr(args, "_vllm_user_provided", set()) - if dest in user_provided: - return True - entry = get_vllm_cli_action_table().get(dest) - if entry is None: - return False - _, action = entry - return getattr(args, dest, action.default) != action.default - - -def _apply_vllm_overrides(args, server_args: dict[str, Any], vllm_overrides: dict | None, rank: int) -> None: - """Merge per-group ``overrides`` from rollout YAML into ``args`` / ``server_args``.""" - if not vllm_overrides: - return - for key, value in vllm_overrides.items(): - normalized = key.replace("-", "_") - if normalized == "model_path": - server_args["model_path"] = value - continue - if normalized.startswith("disaggregation"): - logger.debug("vllm_overrides: skipping unsupported key %s (rank=%s)", key, rank) - continue - dest = normalized if normalized.startswith("vllm_") else f"vllm_{normalized}" - if hasattr(args, dest): - logger.info("vllm_overrides: %s=%r (rank=%s)", dest, value, rank) - setattr(args, dest, value) - continue - if normalized in server_args: - server_args[normalized] = value - continue - logger.debug("vllm_overrides: unrecognized key %s (rank=%s)", key, rank) - - -class _RobustJsonEncoder: - @staticmethod - def default(obj): - import enum - from pathlib import Path - - if dataclasses.is_dataclass(obj): - return dataclasses.asdict(obj) - if isinstance(obj, (set, frozenset)): - return list(obj) - if isinstance(obj, enum.Enum): - return obj.value - if isinstance(obj, Path): - return str(obj) - if isinstance(obj, bytes): - return obj.decode("utf-8", errors="replace") - raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") - - -def serialize_for_cli(value) -> str | None: - import json - - if isinstance(value, str): - return value - if isinstance(value, (int, float, bool)): - return str(value) - if dataclasses.is_dataclass(value) or isinstance(value, (dict, list, tuple)): - try: - return json.dumps(value, default=_RobustJsonEncoder.default) - except (TypeError, ValueError) as exc: - logger.debug("JSON serialization failed for %r: %s", type(value).__name__, exc) - return None - return None - - -def _serialize_weight_transfer_config(value) -> str: - serialized = serialize_for_cli(value) - if serialized is None: - import json - - return json.dumps({"backend": str(value)}) - return serialized - - -def _forward_vllm_cli_args(args, cmd: list[str]) -> None: - """Append user ``--vllm-*`` overrides not already set by the orchestrator.""" - fixed = {flag for flag in cmd if isinstance(flag, str) and flag.startswith("--")} - raw_values: dict[str, str] = getattr(args, "_vllm_raw_values", {}) - - for vime_dest, (vllm_flag, action) in get_vllm_cli_action_table().items(): - if vime_dest in SKIPPED_DESTS: - continue - if vllm_flag in fixed: - continue - if not hasattr(args, vime_dest): - continue - value = getattr(args, vime_dest) - default = action.default - if value == default or value is None: - continue - - if isinstance(action, BooleanOptionalAction): - cmd.append(vllm_flag if value else f"--no-{vllm_flag[2:]}") - continue - if action.nargs == 0: - cmd.append(vllm_flag) - continue - if action.nargs == "?" and action.const is not None and value == action.const: - cmd.append(vllm_flag) - continue - if action.nargs in ("+", "*") or (action.nargs not in (None, "?") and isinstance(value, (list, tuple))): - if not isinstance(value, (list, tuple)): - value = [value] - if not value: - continue - if not all(isinstance(v, _PRIMITIVE_TYPES) for v in value): - logger.debug("Skipping %s: list contains non-primitive items (%r)", vllm_flag, value) - continue - cmd.append(vllm_flag) - cmd.extend(str(v) for v in value) - continue - if not isinstance(value, _PRIMITIVE_TYPES): - raw = raw_values.get(vime_dest) - if raw is not None: - cmd.extend([vllm_flag, raw]) - continue - serialized = serialize_for_cli(value) - if serialized is None: - logger.debug( - "Skipping forward of %s: parsed value %r (%s) cannot be serialized", - vllm_flag, - value, - type(value).__name__, - ) - continue - cmd.extend([vllm_flag, serialized]) - - -def redact_cmd_for_log(cmd: list[str]) -> str: - """Stringify ``cmd`` for logging, redacting credential flags.""" - parts: list[str] = [] - redact_next = False - for token in cmd: - if redact_next: - parts.append("***") - redact_next = False - continue - parts.append(token) - if isinstance(token, str) and token in _REDACTED_FLAGS: - redact_next = True - return " ".join(parts) + return p -def build_vllm_subprocess_env(server_args: dict[str, Any]) -> dict[str, str]: - """Child-process environment for ``vllm serve``.""" - args = server_args["args"] +def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]: + args = server_args_dict["_args"] env = os.environ.copy() env.pop("PYTORCH_CUDA_ALLOC_CONF", None) + env.pop("PYTORCH_ALLOC_CONF", None) env.setdefault("NCCL_CUMEM_ENABLE", "0") - if is_npu(): - env["ASCEND_RT_VISIBLE_DEVICES"] = server_args["visible_devices"] - env["VLLM_USE_AOT_COMPILE"] = "0" - cann_python_path = get_cann_python_site_packages() - if cann_python_path is not None: - prepend_pythonpath(env, cann_python_path) - if getattr(args, "colocate", False): - env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" - else: - env["CUDA_VISIBLE_DEVICES"] = server_args["visible_devices"] + env["CUDA_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"] + # ROCm: keep HIP visibility in sync with CUDA (no-op on CUDA). + env["HIP_VISIBLE_DEVICES"] = server_args_dict["_visible_devices"] + env = current_platform().vllm.subprocess_env( + env, + visible_devices=server_args_dict["_visible_devices"], + colocate=getattr(args, "colocate", False), + ) env.setdefault("VLLM_SERVER_DEV_MODE", "1") + env["VLLM_USE_V2_MODEL_RUNNER"] = "1" if getattr(args, "vllm_enable_deterministic_inference", False): env["VLLM_BATCH_INVARIANT"] = "1" if getattr(args, "colocate", False): @@ -380,178 +85,52 @@ def build_vllm_subprocess_env(server_args: dict[str, Any]) -> dict[str, str]: if vime_root not in {p for p in existing_pp.split(os.pathsep) if p}: env["PYTHONPATH"] = os.pathsep.join(filter(None, [vime_root, existing_pp])) env.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - return env - -def build_vllm_cmd_and_env(server_args: dict[str, Any]) -> tuple[list[str], dict[str, str]]: - """Translate ``server_args`` to ``vllm serve`` argv and subprocess environment.""" - args = server_args["args"] - topology: VllmEngineTopology = server_args["topology"] - env = build_vllm_subprocess_env(server_args) - host_for_subprocess = (server_args["host"] or "127.0.0.1").strip("[]") - - cmd = [ - "vllm", - "serve", - str(server_args["model_path"]), - "--tensor-parallel-size", - str(server_args["tp_size"]), - "--port", - str(server_args["port"]), - "--host", - host_for_subprocess, - "--seed", - str(server_args["seed"]), - "--trust-remote-code", - ] - - if server_args["pp_size"] > 1: - cmd += ["--pipeline-parallel-size", str(server_args["pp_size"])] - - if topology.multi_node: - if server_args["master_addr"] is None or server_args["master_port"] is None: - raise ValueError("master_addr/master_port required for multi-node vLLM engine") - append_vllm_distributed_launch_flags( - cmd, - topology, - (server_args["master_addr"], server_args["master_port"]), - args, - ) - - if getattr(args, "fp16", False): - cmd += ["--dtype", "float16"] - - if getattr(args, "offload_rollout", False) and not getattr(args, "vllm_enable_sleep_mode", False): - cmd += ["--enable-sleep-mode"] - args.vllm_enable_sleep_mode = True - - if ( - getattr(args, "rollout_max_context_len", None) is not None - and getattr(args, "vllm_max_model_len", None) is None - ): - cmd += ["--max-model-len", str(args.rollout_max_context_len)] - - if getattr(args, "use_rollout_routing_replay", False): - cmd += ["--enable-return-routed-experts"] - - # gpu_memory_utilization: no vime-forced default. In colocate, training and rollout do not - # occupy the GPU simultaneously (sleep/offload cycles), so vLLM's own default is fine. A user - # value passed via --vllm-gpu-memory-utilization is auto-forwarded by _forward_vllm_cli_args. - - # 2) logprobs_mode: vllm's raw_logprobs are pre-temperature, while Megatron - # replay compares against rollout-temperature-scaled logprobs. - if not _user_overrode(args, "vllm_logprobs_mode"): - cmd += ["--logprobs-mode", "processed_logprobs"] - - # 3) weight_transfer_config: vllm default None disables /init_weight_transfer_engine, - # so vime's weight sync would fail. - # - Colocated mode: use IPC backend. UpdateWeightFromTensor calls - # IPCWeightTransferEngine.trainer_send_weights and passes an empty init_info - # dict, which is the correct signature for the IPC backend. - # - Non-colocated mode: use NCCL backend. Weight sync goes through - # update_weights_from_distributed; the vLLM engine still needs - # init_weight_transfer_engine to succeed (with NCCL the caller must supply - # master_address, master_port, rank_offset, and world_size separately). - # Users who pass ``--vllm-weight-transfer-config`` explicitly are honored. - if _user_overrode(args, "vllm_weight_transfer_config"): - cmd += [ - "--weight-transfer-config", - _serialize_weight_transfer_config(args.vllm_weight_transfer_config), - ] - elif getattr(args, "colocate", False): - cmd += ["--weight-transfer-config", '{"backend":"ipc"}'] - else: - cmd += ["--weight-transfer-config", '{"backend":"nccl"}'] - - if "--worker-extension-cls" not in cmd: - if getattr(args, "colocate", False): - _ext_cls = ( - "vime.backends.megatron_utils.update_weight.update_weight_from_tensor.vLLMColocateWorkerExtension" - ) - else: - _ext_cls = "vime.backends.megatron_utils.update_weight.update_weight_from_tensor.vLLMWorkerExtension" - cmd += [ - "--worker-extension-cls", - _ext_cls, - ] - - worker_type = server_args.get("worker_type", "regular") - if worker_type in ("prefill", "decode") and topology.node_rank == 0: + worker_type = server_args_dict.get("_worker_type", "regular") + if worker_type in ("prefill", "decode") and server_args_dict.get("node_rank", 0) == 0: + host_for_subprocess = (server_args_dict.get("host") or "127.0.0.1").strip("[]") env["VLLM_NIXL_SIDE_CHANNEL_HOST"] = host_for_subprocess - env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(server_args["disaggregation_bootstrap_port"]) - - _forward_vllm_cli_args(args, cmd) - logger.info("Launching vLLM server: %s", redact_cmd_for_log(cmd)) - return cmd, env - - -def _exec_vllm_cmd(cmd: list[str], env: dict[str, str]) -> None: - """Entry point for multiprocessing child process.""" - os.execvpe(cmd[0], cmd, env) - - -def _normalize_vllm_wake_tags(tags: list[str] | None) -> list[str] | None: - if not tags: - return tags - normalized = [t for t in tags if t in _VLLM_WAKE_TAGS] - dropped = set(tags) - set(normalized) - if dropped: - logger.debug("vLLM wake_up: dropped tags not supported by vLLM: %s", sorted(dropped)) - return normalized or None + env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(server_args_dict["_disaggregation_bootstrap_port"]) + return env -def launch_server_process(server_args: dict) -> multiprocessing.Process: - """Spawn ``vllm serve`` from a :func:`_compute_server_args` dict.""" - cmd, env = build_vllm_cmd_and_env(server_args) - p = _spawn_ctx.Process(target=_exec_vllm_cmd, args=(cmd, env)) - p.start() - return p +def _run_vllm_server(kwargs: dict, env: dict) -> None: + os.environ.update(env) -def _wait_worker_process_alive(process: multiprocessing.Process, timeout_s: float = 300.0) -> None: - """Non-head nodes have no HTTP health endpoint; ensure the subprocess stays up.""" - start = time.time() - while process.is_alive(): - if time.time() - start > timeout_s: - return - time.sleep(2) - raise RuntimeError(f"vLLM worker process exited unexpectedly with code {process.exitcode}") + from vllm.entrypoints.cli.serve import ServeSubcommand + from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args + from vllm.utils.argparse_utils import FlexibleArgumentParser + ns = argparse.Namespace(**kwargs) + parser = make_arg_parser(FlexibleArgumentParser()) + args = parser.parse_args(args=[], namespace=ns) + validate_parsed_serve_args(args) + ServeSubcommand.cmd(args) -def _wait_server_healthy(base_url: str, process: multiprocessing.Process | None) -> None: - """Wait until the vLLM server responds on ``GET /health`` (no time limit, SGLang-style). - Loops until /health returns 200, or — for a managed subprocess — until it dies (fail fast via - ``process.is_alive()``). There is no overall deadline, so a slow-but-healthy startup (a large - MoE / DP engine loading + compiling + capturing CUDA graphs across replicas) is never - spuriously timed out. The per-probe ``timeout=3`` bounds each individual request so a single - stuck socket cannot wedge the loop. In external mode (``process is None``) there is no liveness - signal, so a permanently unreachable URL loops indefinitely by design (the external engine is - caller-managed). Mirrors slime's SGLang backend _wait_server_healthy. - """ +def _wait_server_healthy(base_url, is_process_alive): while True: try: response = requests.get(f"{base_url}/health") if response.status_code == 200: - return + break except requests.RequestException: pass - if process is not None and not process.is_alive(): - raise RuntimeError(f"vLLM server exited unexpectedly with code {process.exitcode}") + if not is_process_alive(): + raise Exception("Server process terminated unexpectedly.") + time.sleep(2) class VLLMEngine(RayActor): - """Ray actor for vLLM OpenAI HTTP rollout (connect or spawn local ``vllm serve``).""" - def __init__( self, args, rank: int, worker_type: str = "regular", base_gpu_id: int | None = None, - model_path: str | None = None, vllm_overrides: dict | None = None, num_gpus_per_engine: int | None = None, ): @@ -559,17 +138,9 @@ def __init__( self.rank = rank self.worker_type = worker_type self.base_gpu_id = base_gpu_id - self.model_path = model_path or args.hf_checkpoint self.vllm_overrides = vllm_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine - self.process: multiprocessing.Process | None = None self._weight_version: str | None = None - self.node_rank = 0 - self._topology: VllmEngineTopology | None = None - self._server_args: dict | None = None - - def _http_base(self) -> str: - return f"http://{self.server_host}:{self.server_port}" def init( self, @@ -581,345 +152,291 @@ def init( router_ip=None, router_port=None, ): - # ``nccl_port`` is allocated by rollout but unused by vLLM (rendezvous uses ``dist_init_addr``). del nccl_port - gpus_per_engine = self.num_gpus_per_engine or self.args.rollout_num_gpus_per_engine + self.router_ip = _wrap_ipv6(router_ip) if router_ip is not None else None + self.router_port = router_port + host = host or get_host_info()[1] - self._server_args = _compute_server_args( + host = _wrap_ipv6(host) + ip_part, port_part = dist_init_addr.rsplit(":", 1) + dist_init_addr = f"{_wrap_ipv6(ip_part)}:{port_part}" + + server_args_dict, external_engine_need_check_fields = _compute_server_args( self.args, self.rank, dist_init_addr, host, port, - worker_type=self.worker_type, + self.worker_type, + disaggregation_bootstrap_port, base_gpu_id=self.base_gpu_id, vllm_overrides=self.vllm_overrides, - num_gpus_per_engine=gpus_per_engine, - disaggregation_bootstrap_port=disaggregation_bootstrap_port, + num_gpus_per_engine=self.num_gpus_per_engine, ) - self._topology = self._server_args["topology"] - self.node_rank = self._topology.node_rank - # rollout always passes the resolved router (engine.init(router_ip=self.router_ip, ...)) - # and _start_router always returns a real address, so no fallback to args is needed. - self.router_ip = router_ip - self.router_port = router_port - self.server_host = self._server_args["host"] - self.server_port = port - self.disaggregation_bootstrap_port = disaggregation_bootstrap_port - - if self.worker_type != "regular": - logger.warning( - "vLLMEngine: worker_type=%s is not used by current vLLM deployment (treated as regular).", - self.worker_type, - ) + self.node_rank = server_args_dict["node_rank"] + self.server_host = server_args_dict["host"] # with [] if ipv6 + self.server_port = server_args_dict["port"] if self.args.rollout_external: - # Only the HTTP-owning head node (node_rank 0) can hit /health and - # /server_info. Headless workers (node_rank>0) expose no HTTP endpoint, - # so skip the check for them — mirrors the head/worker split in _init_normal. - if self.node_rank == 0: - self._init_external() - else: - logger.info( - "External vLLM headless worker (rank=%s node_rank=%s): skip HTTP health/config " - "check (only node_rank 0 owns HTTP).", - self.rank, - self.node_rank, - ) + self._init_external(server_args_dict, external_engine_need_check_fields=external_engine_need_check_fields) else: - self._init_normal() + self._init_normal(server_args_dict) + + def _init_external(self, expect_server_args, external_engine_need_check_fields): + logger.info(f"Use external vLLM engine (rank={self.rank}, expect_server_args={expect_server_args})") + + def _matches_expected(actual, expected): + if isinstance(actual, dict) and isinstance(expected, dict): + return all(key in actual and _matches_expected(actual[key], value) for key, value in expected.items()) + return actual == expected + + def _sanity_check_server_args(actual_server_args, expect_server_args): + for name in external_engine_need_check_fields: + expect_value = expect_server_args.get(name) + actual_value = actual_server_args.get(name) + assert _matches_expected( + actual_value, expect_value + ), f"{name=} {expect_value=} {actual_value=} {expect_server_args=} {actual_server_args=}" + + actual_server_args = get_server_info(f"http://{self.server_host}:{self.server_port}") + _sanity_check_server_args(actual_server_args, expect_server_args) + self._register_to_router(expect_server_args) + + def _init_normal(self, server_args_dict): + logger.info(f"Launch vLLM api_server at: {self.server_host}:{self.server_port}") + self.process = launch_server_process(server_args_dict) + self._register_to_router(server_args_dict) + + def _register_to_router(self, server_args_dict): + if self.worker_type == "encoder": + return if self.node_rank == 0 and self.router_ip and self.router_port: - self._register_worker_with_router() - - def _register_worker_with_router(self) -> None: - worker_url = self._http_base() - payload = {"url": worker_url, "worker_type": self.worker_type} - response = requests.post( - f"http://{self.router_ip}:{self.router_port}/workers", - json=payload, - timeout=30, - ) - response.raise_for_status() - - def _deregister_worker_from_router(self) -> None: - if self.node_rank != 0 or not self.router_ip or not self.router_port: - return - worker_url = self._http_base() - try: - all_workers = requests.get(f"http://{self.router_ip}:{self.router_port}/workers", timeout=30).json()[ - "workers" - ] - for worker in all_workers: - if worker["url"] == worker_url: - response = requests.delete( - f"http://{self.router_ip}:{self.router_port}/workers/{quote(worker_url, safe='')}", - timeout=30, + worker_url = f"http://{self.server_host}:{self.server_port}" + payload = { + "url": worker_url, + "worker_type": self.worker_type, + } + if self.worker_type == "prefill": + bootstrap_port = server_args_dict.get("_disaggregation_bootstrap_port") + if bootstrap_port is None: + bootstrap_port = server_args_dict.get("disaggregation_bootstrap_port") + if bootstrap_port is None: + raise RuntimeError( + f"Prefill worker {worker_url} does not have disaggregation_bootstrap_port; cannot register it to the PD router." ) - response.raise_for_status() - return - logger.warning("Worker %s not found in vllm-router during shutdown.", worker_url) - except Exception as e: - logger.warning("Failed to list/remove worker on vllm-router: %s", e) - - def _init_external(self) -> None: - logger.info("Use external vLLM engine (rank=%s) at %s:%s", self.rank, self.server_host, self.server_port) - _wait_server_healthy(self._http_base(), process=None) - self._sanity_check_external_server_args() - - def _sanity_check_external_server_args(self) -> None: - """Strictly verify an external engine's parallel config matches what we expect; raise on mismatch. - - Replaces the previous warn-only check, which (a) compared against the *global* - ``rollout_num_gpus_per_engine`` — wrong for heterogeneous / multi-node groups — and - (b) only logged a warning, so a misconfigured external engine sailed through and then - hung the weight-sync rendezvous ~300s later with no clear error. - - We now compare every field in ``EXTERNAL_ENGINE_CHECK_FIELDS`` against the per-engine - expectation in ``self._server_args`` and raise immediately on mismatch. A field that the - engine's ``/server_info`` does not report (``actual is None``) is skipped rather than - treated as a mismatch (e.g. vLLM ``parallel_config`` may not surface ``nnodes``), so the - check stays strict for reported fields without false-failing on unreported ones. - """ - response = requests.get(f"{self._http_base()}/server_info", params={"config_format": "json"}, timeout=30) - body = _response_json(response) - parallel_cfg = body.get("vllm_config", {}).get("parallel_config", {}) - if not parallel_cfg: - raise RuntimeError(f"External vLLM /server_info missing vllm_config.parallel_config: {body}") - actual = { - "tp_size": parallel_cfg.get("tensor_parallel_size"), - "pp_size": parallel_cfg.get("pipeline_parallel_size"), - "dp_size": parallel_cfg.get("data_parallel_size"), - "nnodes": parallel_cfg.get("nnodes"), - } - expect = {name: self._server_args.get(name) for name in EXTERNAL_ENGINE_CHECK_FIELDS} - for name in EXTERNAL_ENGINE_CHECK_FIELDS: - actual_value = actual.get(name) - if actual_value is None: - logger.debug("External vLLM /server_info did not report %s; skipping that check.", name) - continue - if actual_value != expect.get(name): - raise AssertionError( - f"External vLLM server arg mismatch: {name}: expect={expect.get(name)} " - f"actual={actual_value} (full expect={expect} actual={actual})" - ) + payload["bootstrap_port"] = bootstrap_port + response = requests.post( + f"http://{self.router_ip}:{self.router_port}/workers", + json=payload, + ) + response.raise_for_status() - def _init_normal(self) -> None: - topology = self._topology - assert topology is not None and self._server_args is not None - logger.info( - "Launch vLLM OpenAI api_server at: %s:%s (rank=%s node_rank=%s/%s)", - self.server_host, - self.server_port, - self.rank, - topology.node_rank, - topology.nnodes, - ) - self.process = launch_server_process(self._server_args) - if topology.node_rank == 0: - _wait_server_healthy(self._http_base(), process=self.process) - else: - _wait_worker_process_alive(self.process) + def _make_request(self, endpoint: str, payload: dict | None = None): + """Make a POST request to the specified endpoint with the given payload. - def _make_request(self, endpoint: str, payload: dict | None = None) -> dict | None: - """Control-plane POST returning parsed JSON (mirrors SGLang's ``_make_request``). + Args: + endpoint: The API endpoint to call + payload: The JSON payload to send (default: empty dict) - The single choke point for control-plane POSTs: headless workers (node_rank>0) own no - HTTP server, so they no-op to None; otherwise POST and parse via the shared - ``_response_json`` (also reused by the query-param endpoints /sleep, /wake_up, ...). + Returns: + The JSON response from the server """ if self.node_rank != 0: - return None - url = f"{self._http_base()}/{endpoint.lstrip('/')}" - return _response_json(requests.post(url, json=payload or {})) - - def _post_vllm_update_weights_http(self, update_info: dict) -> dict: - """POST ``/update_weights`` with ``{"update_info": ...}`` (vLLM RLHF control plane). + return - Caller must invoke ``start_weight_update`` / ``finish_weight_update`` around a batch of - ``/update_weights`` calls (see ``UpdateWeightFromTensor`` / ``UpdateWeightFromDistributed``). - """ - return self._make_request( - "update_weights", - {"update_info": update_info}, - ) + url = f"http://{self.server_host}:{self.server_port}/{endpoint}" + response = requests.post(url, json=payload or {}) + try: + response.raise_for_status() + except requests.exceptions.HTTPError as e: + e.add_note(f"{response.text=}") + raise + if not response.content or not response.content.strip(): + return {"ok": True} + return response.json() def health_generate(self, timeout: float = 5.0) -> bool: - """Return True if ``GET /health`` succeeds.""" if self.node_rank != 0: return True - response = requests.get(f"{self._http_base()}/health", timeout=timeout) + + response = requests.get( + f"http://{self.server_host}:{self.server_port}/health", + timeout=timeout, + ) response.raise_for_status() return True - def update_weights(self, request: dict, weight_version: str | None = None) -> dict | None: - """POST native vLLM ``/update_weights`` payloads from transfer engines.""" - if self.node_rank != 0: - return None - - update_info = request["update_info"] if "update_info" in request else request - payload = dict(update_info) - if "ipc_handles" in payload: - payload["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(payload.pop("ipc_handles"))).decode("utf-8") - - response = self._post_vllm_update_weights_http(payload) - if weight_version is not None: - self._weight_version = str(weight_version) - return response + def update_weights(self, update_info: dict | list[dict | None]): + infos = update_info if isinstance(update_info, list) else [update_info] + payload = [] + for info in infos: + if info is None: + payload.append(None) + continue + worker_payload = dict(info) + ipc_handles = worker_payload.pop("ipc_handles", None) + if ipc_handles is not None: + worker_payload["ipc_handles_pickled"] = base64.b64encode(cloudpickle.dumps(ipc_handles)).decode( + "utf-8" + ) + payload.append(worker_payload) + if not isinstance(update_info, list): + payload = payload[0] + return self._make_request("update_weights", {"update_info": payload}) def flush_cache(self): - """Clear prefix cache via ``POST /reset_prefix_cache``.""" if self.node_rank != 0: return params = {"reset_running_requests": False} - requests.post(f"{self._http_base()}/reset_prefix_cache", params=params).raise_for_status() + requests.post( + f"http://{self.server_host}:{self.server_port}/reset_prefix_cache", params=params + ).raise_for_status() def get_url(self): - """Worker HTTP base URL, or ``None`` when ``node_rank != 0``.""" if self.node_rank != 0: return None - return self._http_base() + return f"http://{self.server_host}:{self.server_port}" def shutdown(self): - logger.info("Shutdown vLLM engine %s:%s...", self.server_host, self.server_port) - self._deregister_worker_from_router() if self.args.rollout_external: return - if self.process is None or not self.process.is_alive(): - return - pid = self.process.pid - try: - from vllm.utils.system_utils import kill_process_tree - - kill_process_tree(pid) - except Exception as e: - logger.warning("vLLM kill_process_tree failed (%s); terminate root only.", e) - if self.process.is_alive(): - self.process.terminate() - try: - self.process.join(timeout=15) - except Exception: - pass - if self.process.is_alive(): - self.process.kill() - try: - self.process.join(timeout=30) - except Exception: - pass - self.process = None + logger.info(f"Shutdown engine {self.server_host}:{self.server_port}...") + if self.worker_type != "encoder" and self.node_rank == 0: + worker_url = f"http://{self.server_host}:{self.server_port}" + try: + all_workers = requests.get(f"http://{self.router_ip}:{self.router_port}/workers").json()["workers"] + for worker in all_workers: + if worker["url"] == worker_url: + response = requests.delete( + f"http://{self.router_ip}:{self.router_port}/workers/{quote(worker_url, safe='')}", + ) + response.raise_for_status() + break + else: + logger.warning(f"Worker {worker_url} not found in vllm-router during shutdown.") + except Exception as e: + logger.warning(f"Failed to fetch workers list or remove worker: {e}") - def get_weight_version(self) -> str | None: - """Return the version recorded by the last successful weight transfer. + kill_process_tree(self.process.pid) - Raises ``RuntimeError`` if no weight transfer has recorded a version - yet — we don't fall back to a ``/v1/models`` lookup, which would - return the model path string and never match the trainer's integer - counter (i.e. produce a misleading "mismatch" downstream). - Worker ranks (``node_rank != 0``) short-circuit per the class idiom. - """ + def get_weight_version(self): if self.node_rank != 0: - return None - if self._weight_version is None: - raise RuntimeError( - "VLLMEngine.get_weight_version called before any successful " - "weight transfer recorded a version (update_weights " - "/ update_weights_from_distributed never reached its " - "post-POST version write)." - ) + return + response = requests.get(f"http://{self.server_host}:{self.server_port}/weight_info") + try: + response.raise_for_status() + except requests.exceptions.HTTPError as error: + error.add_note(f"{response.text=}") + raise + weight_version = response.json()["weight_version"] + self._weight_version = None if weight_version is None else str(weight_version) return self._weight_version + def set_weight_version(self, new_version: str): + version = str(new_version) + result = self._make_request("update_weight_version", {"new_version": version}) + self._weight_version = version + return result + def release_memory_occupation(self, level: int = 2): - """Flush prefix cache, then ``POST /sleep?level={level}``.""" - if self.node_rank != 0: - return None self.flush_cache() - response = requests.post( - f"{self._http_base()}/sleep", - params={"level": level}, - ) - return _response_json(response) + response = requests.post(f"http://{self.server_host}:{self.server_port}/sleep", params={"level": level}) + response.raise_for_status() + if not response.content or not response.content.strip(): + return {"ok": True} + return response.json() - def resume_memory_occupation(self, tags: list[str] | None = None): - """``POST /wake_up`` with vLLM-supported wake tags.""" - if self.node_rank != 0: - return None + def resume_memory_occupation(self, tags: list[str] = None): tags = _normalize_vllm_wake_tags(tags) - # vLLM ``POST /wake_up`` uses ``query_params.getlist("tags")``, not JSON. - # Omit params when ``tags`` is empty so the server wakes all tags (see api_router.wake_up). wake_params: list[tuple[str, str]] | None = [("tags", t) for t in tags] if tags else None - response = requests.post( - f"{self._http_base()}/wake_up", - params=wake_params, - timeout=30, - ) - return _response_json(response) + response = requests.post(f"http://{self.server_host}:{self.server_port}/wake_up", params=wake_params) + response.raise_for_status() + if not response.content or not response.content.strip(): + return {"ok": True} + return response.json() + + def check_weights(self, action: str): + del action + return {"ok": True, "supported": False} def init_weight_transfer_engine(self, payload: dict) -> dict: - """``POST /init_weight_transfer_engine`` with a caller-supplied payload (IPC path). + return self._make_request("init_weight_transfer_engine", payload) - For IPC mode the payload is ``{"init_info": {}}``; for NCCL use - ``init_weights_update_group`` which constructs the payload from typed args. - """ - last_error = None - for attempt in range(1, 4): - try: - return self._make_request("init_weight_transfer_engine", payload) - except Exception as e: - last_error = e - if attempt < 3: - logger.warning("init_weight_transfer_engine attempt %s/3 failed: %s", attempt, e) - time.sleep(2 * attempt) - raise RuntimeError(f"vLLM init_weight_transfer_engine failed: {last_error}") from last_error + def start_weight_update(self) -> dict: + return self._make_request("start_weight_update", {}) - def start_weight_update(self, is_checkpoint_format: bool = True) -> dict: - """``POST /start_weight_update`` — signals vLLM to enter IPC weight-update mode.""" - return self._make_request("start_weight_update", {"is_checkpoint_format": is_checkpoint_format}) + def start_draft_weight_update(self) -> dict: + return self._make_request("start_draft_weight_update", {}) - def finish_weight_update(self) -> dict: - """``POST /finish_weight_update`` — signals vLLM to exit IPC weight-update mode. + def finish_weight_update(self, weight_version: str | None = None) -> dict: + payload = {} if weight_version is None else {"weight_version": str(weight_version)} + result = self._make_request("finish_weight_update", payload) + if weight_version is not None: + self._weight_version = str(weight_version) + return result - Purely a state-machine bookend; ``_weight_version`` is recorded by - the data-carrying ``update_weights`` / distributed update calls. - """ - return self._make_request("finish_weight_update", {}) + def pull_weights(self, target_version: int): + if self.node_rank != 0: + return + response = requests.post( + f"http://{self.server_host}:{self.server_port}/collective_rpc", + json={ + "method": "pull_weights", + "kwargs": { + "local_checkpoint_dir": self.args.update_weight_local_checkpoint_dir, + "source_dir": self.args.update_weight_disk_dir, + "target_version": target_version, + "pre_read_hook": self.args.custom_update_weight_pre_read_path, + }, + }, + ) + response.raise_for_status() + result = response.json() + self.set_weight_version(str(target_version)) + return result - def check_weights(self, action: str): - """No vLLM ``weights_checker`` route; return a placeholder dict.""" - del action - return {"ok": True, "supported": False, "note": "vLLM has no weights_checker endpoint."} + def update_weights_from_disk( + self, + model_path: str, + load_format: str | None = None, + weight_version: str | None = None, + ): + del load_format + if self.node_rank != 0: + return + response = requests.post( + f"http://{self.server_host}:{self.server_port}/collective_rpc", + json={"method": "reload_weights", "kwargs": {"weights_path": model_path, "is_checkpoint_format": True}}, + ) + try: + response.raise_for_status() + except requests.exceptions.HTTPError as e: + e.add_note(f"{response.text=}") + raise + if weight_version is not None: + self.set_weight_version(str(weight_version)) + return response.json() def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): - """Call ``POST /init_weight_transfer_engine`` with an ``init_info`` block. - - ``group_name`` / ``backend`` are accepted for a uniform caller signature but are not sent to vLLM. - Always uses the vllm-native weight transfer engine; reload-on-continue fallback is no longer supported. - """ del group_name, backend - payload = { - "init_info": { - "master_address": master_address, - "master_port": master_port, - "rank_offset": rank_offset, - "world_size": world_size, - } - } - last_error = None - for attempt in range(1, 4): - try: - return self._make_request("init_weight_transfer_engine", payload) - except Exception as e: - last_error = e - if attempt < 3: - logger.warning("init_weight_transfer_engine attempt %s/3 failed: %s", attempt, e) - time.sleep(2 * attempt) - raise RuntimeError(f"vLLM init_weight_transfer_engine failed: {last_error}") from last_error + return self._make_request( + "init_weight_transfer_engine", + { + "init_info": { + "master_address": master_address, + "master_port": master_port, + "rank_offset": rank_offset, + "world_size": world_size, + } + }, + ) def destroy_weights_update_group(self, group_name): - """No vLLM destroy call; return ``None``.""" del group_name return None @@ -928,18 +445,10 @@ def update_weights_from_distributed( names, dtypes, shapes, - group_name, + *, flush_cache=False, - weight_version: str | None = None, - packed: bool = True, + weight_version: str, ): - """NCCL path: ``POST /update_weights`` with packed tensor metadata. - - Payload matches vLLM NCCL weight transfer (see upstream rlhf_http_nccl example). - """ - del group_name - if weight_version is not None: - self._weight_version = str(weight_version) if flush_cache: self.flush_cache() dtype_names = [str(d).replace("torch.", "") for d in dtypes] @@ -947,43 +456,27 @@ def update_weights_from_distributed( "names": names, "dtype_names": dtype_names, "shapes": [list(s) for s in shapes], - "packed": bool(packed), + "packed": True, } - return self._post_vllm_update_weights_http(update_info) - - def update_weights_from_disk(self, model_path: str, load_format: str | None = None): - """``POST /collective_rpc`` with ``reload_weights`` and ``weights_path``.""" - if self.node_rank != 0: - return - del load_format - response = requests.post( - f"{self._http_base()}/collective_rpc", - json={ - "method": "reload_weights", - "kwargs": {"weights_path": model_path, "is_checkpoint_format": True}, - }, - timeout=600, - ) - return _response_json(response) + result = self._make_request("update_weights", {"update_info": update_info}) + del weight_version + return result def pause_generation(self): - """``POST /pause`` with mode="keep"; returns the ``requests.Response``.""" if self.node_rank != 0: - return None + return response = requests.post( - f"{self._http_base()}/pause", + f"http://{self.server_host}:{self.server_port}/pause", params={"mode": "keep", "clear_cache": "false"}, json={}, - timeout=120, ) response.raise_for_status() return response def continue_generation(self): - """``POST /resume`` to continue generation after pause.""" if self.node_rank != 0: - return None - response = requests.post(f"{self._http_base()}/resume", json={}, timeout=120) + return + response = requests.post(f"http://{self.server_host}:{self.server_port}/resume", json={}) response.raise_for_status() return response @@ -992,9 +485,8 @@ def post_process_weights( restore_weights_before_load: bool = False, post_process_quantization: bool = False, ): - """No vLLM HTTP hook for post-load processing; return a noop placeholder dict.""" del restore_weights_before_load, post_process_quantization - return {"ok": True, "noop": True, "note": "vLLM post_process is internal to load; no HTTP API."} + return {"ok": True, "noop": True} def start_profile( self, @@ -1006,31 +498,16 @@ def start_profile( with_stack: bool | None = None, record_shapes: bool | None = None, ): - """``POST /start_profile`` with an empty JSON body; kwargs are not forwarded and may be ignored by the server.""" if self.node_rank != 0: - return None - if any( - x is not None and x is not False - for x in ( - output_dir, - start_step, - num_steps, - activities, - profile_by_stage, - with_stack, - record_shapes, - ) - ): - logger.warning("vLLM start_profile: extra kwargs may be ignored by server; see vLLM profiling docs.") - response = requests.post(f"{self._http_base()}/start_profile", json={}, timeout=30) + return + response = requests.post(f"http://{self.server_host}:{self.server_port}/start_profile", json={}) response.raise_for_status() return response def stop_profile(self): - """POST ``/stop_profile`` to stop an active server-side profile.""" if self.node_rank != 0: - return None - response = requests.post(f"{self._http_base()}/stop_profile", json={}, timeout=30) + return + response = requests.post(f"http://{self.server_host}:{self.server_port}/stop_profile", json={}) response.raise_for_status() return response @@ -1041,61 +518,260 @@ def simulate_crash(self): self.args.rollout_external, ) return - logger.info("Simulating crash on vLLM engine %s:%s...", self.server_host, self.server_port) + + logger.info(f"Simulating crash on engine {self.server_host}:{self.server_port}...") self.shutdown() +def _normalize_vllm_wake_tags(tags: list[str] | None) -> list[str] | None: + if not tags: + return tags + normalized = [t for t in tags if t in _VLLM_WAKE_TAGS] + dropped = set(tags) - set(normalized) + if dropped: + logger.debug("vLLM wake_up: dropped tags not supported by vLLM: %s", sorted(dropped)) + return normalized or None + + +def _resolve_parallel_sizes( + args, *, gpus_per_engine: int, overrides: dict[str, Any] | None = None +) -> tuple[int, int, int, int]: + overrides = {key.replace("-", "_"): value for key, value in (overrides or {}).items()} + legacy_fields = _LEGACY_VLLM_PARALLEL_FIELDS.intersection(overrides) + if legacy_fields: + raise ValueError( + "vLLM overrides must use native field names: tensor_parallel_size, " + "pipeline_parallel_size, prefill_context_parallel_size, and data_parallel_size; " + f"got {sorted(legacy_fields)}" + ) + invalid_fields = _INVALID_VLLM_PARALLEL_FIELDS.intersection(overrides) + if invalid_fields: + raise ValueError( + "vLLM does not accept explicit EP/MoE-DP sizes; use " + "enable_expert_parallel with TP/PCP/DP instead: " + f"{sorted(invalid_fields)}" + ) + pp = int(overrides.get("pipeline_parallel_size", getattr(args, "vllm_pipeline_parallel_size", 1)) or 1) + pcp = int( + overrides.get( + "prefill_context_parallel_size", + getattr(args, "vllm_prefill_context_parallel_size", 1), + ) + or 1 + ) + dp = int( + overrides.get( + "data_parallel_size", + getattr(args, "vllm_dp_size", None) or getattr(args, "vllm_data_parallel_size", 1), + ) + or 1 + ) + tp_override = overrides.get("tensor_parallel_size") + parallel_divisor = pp * pcp * dp + if tp_override is None and gpus_per_engine % parallel_divisor != 0: + raise ValueError( + f"num_gpus_per_engine ({gpus_per_engine}) must be divisible by " + "vllm_pipeline_parallel_size * vllm_prefill_context_parallel_size * " + f"vllm_data_parallel_size ({pp} * {pcp} * {dp} = {parallel_divisor})" + ) + tp = int(tp_override) if tp_override is not None else gpus_per_engine // parallel_divisor + if tp * pp * pcp * dp != gpus_per_engine: + raise ValueError( + f"num_gpus_per_engine ({gpus_per_engine}) must equal tensor_parallel_size * " + "pipeline_parallel_size * prefill_context_parallel_size * data_parallel_size " + f"({tp} * {pp} * {pcp} * {dp} = {tp * pp * pcp * dp})" + ) + return tp, pp, pcp, dp + + def _compute_server_args( args, rank, dist_init_addr, host, port, - *, worker_type: str = "regular", + disaggregation_bootstrap_port: int | None = None, base_gpu_id: int | None = None, vllm_overrides: dict | None = None, num_gpus_per_engine: int | None = None, - disaggregation_bootstrap_port: int | None = None, -) -> dict[str, Any]: - """Build per-actor launch config for ``launch_server_process``.""" - gpus_per_engine = num_gpus_per_engine or args.rollout_num_gpus_per_engine - if gpus_per_engine > args.num_gpus_per_node and gpus_per_engine % args.num_gpus_per_node != 0: - raise ValueError( - "vLLM multi-node rollout requires rollout_num_gpus_per_engine to be divisible by " - f"num_gpus_per_node, got rollout_num_gpus_per_engine={gpus_per_engine} " - f"num_gpus_per_node={args.num_gpus_per_node}." - ) +): + normalized_overrides = {} + for key, value in (vllm_overrides or {}).items(): + normalized_key = key.replace("-", "_") + normalized_overrides[normalized_key] = value + vllm_overrides = normalized_overrides + ec_transfer_override = vllm_overrides.pop("ec_transfer_config", None) + + _gpus_per_engine = num_gpus_per_engine or args.rollout_num_gpus_per_engine + nnodes = max(1, _gpus_per_engine // args.num_gpus_per_node) + node_rank = rank % nnodes + if nnodes == 1: + local_num_gpus = min(args.num_gpus_per_node, _gpus_per_engine) + else: + if _gpus_per_engine % nnodes != 0: + raise ValueError( + f"rollout_num_gpus_per_engine ({_gpus_per_engine}) must be divisible by the number of nodes per engine ({nnodes})" + ) + local_num_gpus = _gpus_per_engine // nnodes - topology = compute_vllm_engine_topology(args, rank, num_gpus_per_engine=gpus_per_engine) + tp, pp, pcp, dp = _resolve_parallel_sizes(args, gpus_per_engine=_gpus_per_engine, overrides=vllm_overrides) base = base_gpu_id if base_gpu_id is not None else get_base_gpu_id(args, rank) master_addr: str | None = None master_port: int | None = None - if topology.multi_node: + if nnodes > 1: if not dist_init_addr: raise ValueError("dist_init_addr is required when launching a multi-node vLLM engine") - master_addr, master_port = parse_dist_init_addr(dist_init_addr) - - server_args = { - "args": args, - "rank": rank, - "worker_type": worker_type, - "model_path": args.hf_checkpoint, - "host": _format_v6_uri(host), + ip_part, port_part = dist_init_addr.rsplit(":", 1) + master_addr = ip_part.strip("[]") + master_port = int(port_part) + + kwargs: dict[str, Any] = { + "model": str(args.hf_checkpoint), + "trust_remote_code": True, + "seed": args.seed + rank * args.num_gpus_per_node, + "host": _wrap_ipv6(host or "127.0.0.1"), "port": port, - "master_addr": master_addr, - "master_port": master_port, - "dist_init_addr": dist_init_addr, - "nnodes": topology.nnodes, - "node_rank": topology.node_rank, - "topology": topology, - "visible_devices": ",".join(str(base + i) for i in range(topology.local_num_gpus)), - "tp_size": topology.tensor_parallel_size, - "pp_size": topology.pipeline_parallel_size, - "dp_size": _get_vllm_dp_size(args), - "seed": getattr(args, "seed", 1234) + rank, - "disaggregation_bootstrap_port": disaggregation_bootstrap_port, + "nnodes": nnodes, + "node_rank": node_rank, + "tensor_parallel_size": tp, + "logprobs_mode": "processed_logprobs", + "enable_prompt_tokens_details": True, + "enable_server_load_tracking": True, } - _apply_vllm_overrides(args, server_args, vllm_overrides, rank) - return server_args + + if pp > 1: + kwargs["pipeline_parallel_size"] = pp + + if nnodes > 1: + kwargs["master_addr"] = master_addr + kwargs["master_port"] = master_port + kwargs["data_parallel_backend"] = "mp" + kwargs["distributed_executor_backend"] = "mp" + if node_rank != 0: + kwargs["headless"] = True + + if worker_type == "prefill": + assert ( + disaggregation_bootstrap_port is not None + ), "disaggregation_bootstrap_port must be set for prefill worker" + kwargs["kv_transfer_config"] = { + "kv_connector": "NixlConnector", + "kv_role": "kv_producer", + } + elif worker_type == "decode": + kwargs["kv_transfer_config"] = { + "kv_connector": "NixlConnector", + "kv_role": "kv_consumer", + } + + if ec_transfer_override is not None and worker_type in ("encoder", "regular", "prefill"): + kwargs["ec_transfer_config"] = { + "ec_connector": "ECExampleConnector", + "ec_role": "ec_producer" if worker_type == "encoder" else "ec_consumer", + **ec_transfer_override, + } + + if args.use_rollout_routing_replay: + kwargs["enable_return_routed_experts"] = True + if getattr(args, "vllm_speculative_config", None) is not None: + kwargs["per_request_spec_decode_metrics"] = "summary" + if getattr(args, "rollout_top_p", 1.0) != 1.0: + kwargs["return_sampling_mask"] = True + if args.fp16: + kwargs["dtype"] = "float16" + + if args.offload_rollout and not getattr(args, "vllm_enable_sleep_mode", False): + kwargs["enable_sleep_mode"] = True + args.vllm_enable_sleep_mode = True + + if ( + getattr(args, "rollout_max_context_len", None) is not None + and getattr(args, "vllm_max_model_len", None) is None + ): + kwargs["max_model_len"] = args.rollout_max_context_len + + if args.colocate: + kwargs["weight_transfer_config"] = {"backend": "ipc"} + else: + kwargs["weight_transfer_config"] = {"backend": "nccl"} + kwargs["weight_transfer_config"]["backend"] = current_platform().weight_transfer.backend( + kwargs["weight_transfer_config"]["backend"] + ) + + if worker_type == "encoder": + # vLLM EPD producers have no language-model KV cache groups. Prefix + # caching must therefore be disabled; vLLM's EPD reference launcher + # uses the same setting. + kwargs["enable_prefix_caching"] = False + + external_engine_need_check_fields = [k for k in kwargs.keys() if k not in _EXTERNAL_ENGINE_SKIP_CHECK_FIELDS] + + global _VLLM_SERVER_FIELDS # noqa: PLW0603 + if _VLLM_SERVER_FIELDS is None: + _VLLM_SERVER_FIELDS = _vllm_server_field_names() + + for key, value in vars(args).items(): + if not key.startswith("vllm_"): + continue + field_name = key[len("vllm_") :] + if field_name not in _VLLM_SERVER_FIELDS: + continue + if field_name in kwargs: + continue + if value is None: + continue + kwargs[field_name] = value + + # Per-server-group overrides from --vllm-config YAML. + # Applied after base args so they take highest priority. + if vllm_overrides: + for key, value in vllm_overrides.items(): + if key in ("model_path",) or key.startswith("disaggregation"): + continue + if key in kwargs: + logger.info(f"vllm_overrides: overriding {key}={kwargs[key]} -> {value} (rank={rank})") + kwargs[key] = value + if "model_path" in vllm_overrides: + kwargs["model"] = str(vllm_overrides["model_path"]) + + kwargs["host"] = _wrap_ipv6(kwargs.get("host") or "127.0.0.1") + + # vLLM-specific: topology metadata consumed by launch_server_process / _build_subprocess_env. + # These keys are stripped before passing to vLLM's argparse. + kwargs["_args"] = args + kwargs["_rank"] = rank + kwargs["_worker_type"] = worker_type + kwargs["_visible_devices"] = ",".join(str(base + i) for i in range(local_num_gpus)) + kwargs["_tp_size"] = tp + kwargs["_pp_size"] = pp + kwargs["_pcp_size"] = pcp + kwargs["_dp_size"] = dp + kwargs["_disaggregation_bootstrap_port"] = disaggregation_bootstrap_port + + return kwargs, external_engine_need_check_fields + + +def _vllm_server_field_names() -> frozenset[str]: + """Return the vLLM fields accepted by CLI generation and config overrides.""" + from vllm.engine.arg_utils import AsyncEngineArgs + from vllm.entrypoints.launchers.cli_args import FrontendArgs + + return frozenset(f.name for f in (*dataclasses.fields(AsyncEngineArgs), *dataclasses.fields(FrontendArgs))) + + +_VLLM_SERVER_FIELDS: frozenset[str] | None = None + + +_EXTERNAL_ENGINE_SKIP_CHECK_FIELDS = [ + "model", + "trust_remote_code", + "seed", + "host", + "port", + "tensor_parallel_size", + "logprobs_mode", + "enable_prompt_tokens_details", + "enable_server_load_tracking", +] diff --git a/vime/observability/__init__.py b/vime/observability/__init__.py new file mode 100644 index 000000000..3dab47922 --- /dev/null +++ b/vime/observability/__init__.py @@ -0,0 +1 @@ +"""Training and rollout observability utilities.""" diff --git a/vime/utils/logging_utils.py b/vime/observability/logging_utils.py similarity index 84% rename from vime/utils/logging_utils.py rename to vime/observability/logging_utils.py index 11348a407..428b798f1 100644 --- a/vime/utils/logging_utils.py +++ b/vime/observability/logging_utils.py @@ -2,13 +2,13 @@ import wandb -from . import wandb_utils -from .tensorboard_utils import _TensorboardAdapter +from vime.observability import wandb_utils +from vime.observability.tensorboard_utils import _TensorboardAdapter _LOGGER_CONFIGURED = False -# ref: SGLang +# ref: vLLM def configure_logger(prefix: str = ""): global _LOGGER_CONFIGURED if _LOGGER_CONFIGURED: @@ -31,10 +31,6 @@ def init_tracking(args, primary: bool = True, **kwargs): wandb_utils.init_wandb_secondary(args, **kwargs) -def update_tracking_open_metrics(args, router_addr): - wandb_utils.reinit_wandb_primary_with_open_metrics(args, router_addr) - - def finish_tracking(args): if not args.use_wandb: return diff --git a/vime/utils/metric_utils.py b/vime/observability/metric_utils.py similarity index 100% rename from vime/utils/metric_utils.py rename to vime/observability/metric_utils.py diff --git a/vime/utils/profile_utils.py b/vime/observability/profile_utils.py similarity index 65% rename from vime/utils/profile_utils.py rename to vime/observability/profile_utils.py index 2aaa4660a..fc015dde9 100644 --- a/vime/utils/profile_utils.py +++ b/vime/observability/profile_utils.py @@ -5,6 +5,7 @@ import torch +from vime.utils import accelerator from vime.utils.memory_utils import print_memory logger = logging.getLogger(__name__) @@ -16,10 +17,10 @@ def __init__(self, args): self._torch_profiler_overall = None self._memory_profiler_overall = None - if args.use_pytorch_profiler and ("train_overall" in args.profile_target): + if args.use_pytorch_profiler: self._torch_profiler_overall = _create_torch_profiler(args, name="train_overall") - if args.record_memory_history and ("train_overall" in args.profile_target): + if args.record_memory_history: self._memory_profiler_overall = _BaseMemoryProfiler.create(args) self._memory_profiler_overall.start() @@ -38,29 +39,16 @@ def step(self, rollout_id: int): ): self._memory_profiler_overall.stop() - def iterate_train_actor(self, iterator): - return _profile_simple_loop(iterator, self.args, name="train_actor") - - def iterate_train_log_probs(self, iterator): - return _profile_simple_loop(iterator, self.args, name="train_log_probs") - - -def _profile_simple_loop(iterator, args, name): - if not (args.use_pytorch_profiler and (name in args.profile_target)): - yield from iterator - return - - torch_profiler = _create_torch_profiler(args, name=name) - torch_profiler.start() - for item in iterator: - yield item - torch_profiler.step() - def _create_torch_profiler(args, name): + activities = [torch.profiler.ProfilerActivity.CPU] + activity_name = accelerator.device_type().upper() + if hasattr(torch.profiler.ProfilerActivity, activity_name): + activities.append(getattr(torch.profiler.ProfilerActivity, activity_name)) + return torch.profiler.profile( + activities=activities, schedule=torch.profiler.schedule( - # TODO the train_actor and train_log_probs ones may need to have different args to control step wait=max(args.profile_step_start - 1, 0), warmup=1 if args.profile_step_start > 0 else 0, active=args.profile_step_end - args.profile_step_start, @@ -101,30 +89,55 @@ def stop(self): class _TorchMemoryProfiler(_BaseMemoryProfiler): + def __init__(self, args): + super().__init__(args) + self._recording = False + + @staticmethod + def _memory_module(): + return accelerator.memory_module() + def start(self): logger.info("Attach OOM dump memory history.") - - torch.cuda.memory._record_memory_history( + memory_module = self._memory_module() + if memory_module is None or not hasattr(memory_module, "_record_memory_history"): + logger.warning("Accelerator memory history is unavailable; skip torch memory profiler.") + return + if not hasattr(memory_module, "_dump_snapshot"): + logger.warning("Accelerator memory snapshot is unavailable; skip torch memory profiler.") + return + + memory_module._record_memory_history( max_entries=1000000, - # record stack information for the trace events - # trace_alloc_record_context=True, stacks="all", ) + self._recording = True def oom_observer(device, alloc, device_alloc, device_free): logger.info( f"Observe OOM, will dump snapshot to {self._path_dump}. ({device=} {alloc=} {device_alloc=} {device_free=}; stacktrace is as follows)" ) traceback.print_stack() - torch.cuda.memory._dump_snapshot(self._path_dump) + memory_module._dump_snapshot(str(self._path_dump)) print_memory("when oom") - torch._C._cuda_attach_out_of_memory_observer(oom_observer) + attach_oom_observer = getattr(torch._C, "_cuda_attach_out_of_memory_observer", None) + if attach_oom_observer is not None: + attach_oom_observer(oom_observer) + else: + logger.warning("Accelerator OOM observer is unavailable; memory snapshot on OOM is disabled.") def stop(self): + if not self._recording: + return logger.info(f"Dump memory snapshot to: {self._path_dump}") - torch.cuda.memory._dump_snapshot(self._path_dump) - torch.cuda.memory._record_memory_history(enabled=None) + memory_module = self._memory_module() + if memory_module is None or not hasattr(memory_module, "_dump_snapshot"): + logger.warning("Accelerator memory snapshot is unavailable; skip dump.") + return + memory_module._dump_snapshot(str(self._path_dump)) + memory_module._record_memory_history(enabled=None) + self._recording = False class _MemrayMemoryProfiler(_BaseMemoryProfiler): diff --git a/vime/observability/rollout_data_utils.py b/vime/observability/rollout_data_utils.py new file mode 100644 index 000000000..ab69784ef --- /dev/null +++ b/vime/observability/rollout_data_utils.py @@ -0,0 +1,153 @@ +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from vime.utils.types import Sample + +logger = logging.getLogger(__name__) + +_ROLLOUT_DATA_TENSOR_DTYPES = { + "tokens": torch.long, + "loss_masks": torch.int, + "rollout_log_probs": torch.float32, + "rollout_top_p_token_ids": torch.int32, + "rollout_top_p_token_offsets": torch.int32, + "teacher_log_probs": torch.float32, + "rollout_routed_experts": None, +} + + +def _cpu_tensor(value, dtype: torch.dtype | None = None) -> torch.Tensor: + if isinstance(value, np.ndarray) and not value.flags.writeable: + value = value.copy() + tensor = torch.as_tensor(value, dtype=dtype) if dtype is not None else torch.as_tensor(value) + return tensor.detach().cpu().contiguous() + + +def tensorize_rollout_data_for_training(rollout_data: dict[str, Any]) -> None: + for key, dtype in _ROLLOUT_DATA_TENSOR_DTYPES.items(): + if key in rollout_data: + rollout_data[key] = [_cpu_tensor(value, dtype=dtype) for value in rollout_data[key]] + + if "multimodal_train_inputs" in rollout_data: + rollout_data["multimodal_train_inputs"] = [ + ( + { + key: _cpu_tensor(value) if isinstance(value, (np.ndarray, torch.Tensor)) else value + for key, value in mm_dict.items() + } + if mm_dict is not None + else None + ) + for mm_dict in rollout_data["multimodal_train_inputs"] + ] + + if "rollout_mask_sums" in rollout_data: + rollout_data["rollout_mask_sums"] = _cpu_tensor( + rollout_data["rollout_mask_sums"], + dtype=torch.float32, + ) + + +def validate_rollout_routed_experts_for_replay( + routed_experts: list[torch.Tensor], + args, +) -> None: + """Reject incomplete PP routing captures before R3 consumes them.""" + if not routed_experts: + raise ValueError("R3 is enabled but no rollout routed-experts tensors were returned.") + + num_layers = int(args.num_layers) + topk = int(args.moe_router_topk) + moe_layer_freq = getattr(args, "moe_layer_freq", None) + if isinstance(moe_layer_freq, (list, tuple)): + moe_layers = [layer_id for layer_id, freq in enumerate(moe_layer_freq[:num_layers]) if int(freq) != 0] + else: + moe_layers = list(range(num_layers)) + + for sample_idx, experts in enumerate(routed_experts): + experts = torch.as_tensor(experts) + if experts.ndim != 3 or tuple(experts.shape[1:]) != (num_layers, topk): + raise ValueError( + "Invalid rollout routed-experts shape for R3: " + f"sample={sample_idx}, got={tuple(experts.shape)}, " + f"expected=(*, {num_layers}, {topk})." + ) + if experts.shape[0] == 0: + raise ValueError(f"R3 sample {sample_idx} has no routed-experts rows.") + if topk > 1: + missing_layers = [layer_id for layer_id in moe_layers if not torch.count_nonzero(experts[:, layer_id, :])] + if missing_layers: + raise ValueError( + "R3 routed-experts capture is all zero for MoE layers " + f"{missing_layers} in sample {sample_idx}. This usually means " + "vLLM pipeline stages did not aggregate their disjoint routing " + "captures; refusing to replay expert 0 everywhere." + ) + + +def validate_rollout_id_annotated(node, depth=0): + """Walk the rollout function's nested output and validate ``rollout_id`` only + when a compact / subagent pattern is detected. + + "Compact" = the rollout function wraps multiple training samples from one + rollout execution into a ``list[Sample]``. In vime's convention the + default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout) + so its leaf ``list[Sample]`` lands at depth 1 and we skip validation, + preserving backward compatibility. A compact rollout adds a third level: + ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout), + so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require + every sibling to carry a non-None ``rollout_id`` and to share the same + value, so the loss reducer counts the rollout once instead of N times. + """ + if isinstance(node, Sample): + return + assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}" + if node and isinstance(node[0], Sample): + if depth >= 2 and len(node) > 1: + rids = [sample.rollout_id for sample in node] + missing = [i for i, rollout_id in enumerate(rids) if rollout_id is None] + assert not missing, ( + f"Compact rollout returned {len(node)} samples but rollout_id is unset on " + f"positions {missing}. Set Sample.rollout_id on every sibling so the loss " + "reducer can aggregate them as one rollout instead of N." + ) + assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}." + return + for item in node: + validate_rollout_id_annotated(item, depth + 1) + + +def load_debug_rollout_data(path_template, *, rollout_id: int, subsample_ratio=None) -> list[Sample]: + data = torch.load(path_template.format(rollout_id=rollout_id), weights_only=False)["samples"] + data = [Sample.from_dict(sample) for sample in data] + if subsample_ratio is not None: + original_num_rows = len(data) + rough_subsample_num_rows = int(original_num_rows * subsample_ratio) + data = data[: rough_subsample_num_rows // 2] + data[-rough_subsample_num_rows // 2 :] + logger.info( + "Subsample loaded debug rollout data using ratio=%s and change num rows %s -> %s", + subsample_ratio, + original_num_rows, + len(data), + ) + return data + + +def save_debug_rollout_data(path_template, data, *, rollout_id: int, evaluation: bool) -> None: + if path_template is None: + return + + path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id))) + logger.info(f"Save debug rollout data to {path}") + path.parent.mkdir(parents=True, exist_ok=True) + + if evaluation: + samples = [sample.to_dict() for info in data.values() for sample in info["samples"]] + else: + samples = [sample.to_dict() for sample in data] + + torch.save({"rollout_id": rollout_id, "samples": samples}, path) diff --git a/vime/observability/rollout_metrics.py b/vime/observability/rollout_metrics.py new file mode 100644 index 000000000..0e31bdfec --- /dev/null +++ b/vime/observability/rollout_metrics.py @@ -0,0 +1,271 @@ +import logging +from typing import Any + +import numpy as np +import torch + +from vime.observability import logging_utils +from vime.observability.metric_utils import ( + compute_pass_rate, + compute_rollout_step, + compute_statistics, + dict_add_prefix, + has_repetition, +) +from vime.utils.misc import group_by, load_function +from vime.utils.types import Sample + +logger = logging.getLogger(__name__) + +_VLLM_REQUEST_PERF_FIELDS = ( + ("request/e2e_latency", "e2e_latency"), + ("request/queue_time", "queue_time"), + ("decode/throughput", "decode_throughput"), +) +_VLLM_PREFILL_PERF_FIELDS = ( + ("prefill/bootstrap_queue_duration", "pd_prefill_bootstrap_queue_duration"), + ("prefill/bootstrap_duration", "pd_prefill_bootstrap_duration"), + ("prefill/alloc_wait_duration", "pd_prefill_alloc_wait_duration"), + ("prefill/forward_duration", "pd_prefill_forward_duration"), + ("prefill/transfer_queue_duration", "pd_prefill_transfer_queue_duration"), + ("prefill/transfer_speed_gb_s", "pd_transfer_speed_gb_s"), + ("prefill/transfer_total_mb", "pd_transfer_total_mb"), + ("prefill/retry_count", "pd_prefill_retry_count"), +) +_VLLM_DECODE_PERF_FIELDS = ( + ("decode/prealloc_duration", "pd_decode_prealloc_duration"), + ("decode/bootstrap_duration", "pd_decode_bootstrap_duration"), + ("decode/alloc_wait_duration", "pd_decode_alloc_wait_duration"), + ("decode/transfer_duration", "pd_decode_transfer_duration"), + ("decode/forward_duration", "pd_decode_forward_duration"), +) + + +def compute_metrics_from_samples(args, samples): + response_lengths = [sample.effective_response_length for sample in samples] + + log_dict = {} + log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/") + log_dict |= _compute_zero_std_metrics(args, samples) + log_dict |= _compute_spec_metrics(args, samples) + log_dict |= _compute_prefix_cache_metrics(samples) + log_dict |= _compute_reward_cat_metrics(args, samples) + log_dict |= _compute_top_p_kept_vocab_metrics(samples) + log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item() + log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item() + return log_dict + + +def compute_perf_metrics_from_samples(args, samples, rollout_time): + non_generation_time = [sample.non_generation_time for sample in samples] + + log_dict = {} + log_dict["rollout_time"] = rollout_time + if max(non_generation_time) > 0: + log_dict |= dict_add_prefix(compute_statistics(non_generation_time), "non_generation_time/") + + def token_perf(response_lengths, non_generation_time, key=""): + max_response_length = max(response_lengths) + if args.rollout_num_gpus: + log_dict[f"{key}tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus + log_dict[f"longest_{key}sample_tokens_per_sec"] = max_response_length / rollout_time + + if max(non_generation_time) == 0: + return + + non_generation_time = [ + t for t, length in zip(non_generation_time, response_lengths, strict=True) if length == max_response_length + ] + mean_non_generation_time = sum(non_generation_time) / len(non_generation_time) + + log_dict[f"longest_{key}sample_non_generation_time"] = mean_non_generation_time + log_dict[f"longest_{key}sample_tokens_per_sec_without_non_generation"] = max_response_length / ( + rollout_time - mean_non_generation_time + ) + + token_perf([sample.response_length for sample in samples], non_generation_time, key="") + token_perf([sample.effective_response_length for sample in samples], non_generation_time, key="effective_") + log_dict |= _compute_vllm_request_perf_metrics(samples) + + return log_dict + + +def _compute_vllm_request_perf_metrics(all_samples: list[Sample]): + attrs_by_request = list(_iter_vllm_generate_attrs(all_samples)) + if not attrs_by_request: + return {} + + values_by_metric: dict[str, list[float]] = {} + profiled_request_count = 0 + + def add_value(metric_key: str, source_key: str, attrs: dict) -> bool: + value = attrs.get(source_key) + if not isinstance(value, (int, float)) or isinstance(value, bool) or not np.isfinite(value): + return False + values_by_metric.setdefault(metric_key, []).append(float(value)) + return True + + for attrs in attrs_by_request: + request_has_perf = False + + for metric_key, source_key in _VLLM_REQUEST_PERF_FIELDS: + request_has_perf |= add_value(metric_key, source_key, attrs) + + for metric_key, source_key in _VLLM_PREFILL_PERF_FIELDS: + request_has_perf |= add_value(metric_key, source_key, attrs) + + for metric_key, source_key in _VLLM_DECODE_PERF_FIELDS: + request_has_perf |= add_value(metric_key, source_key, attrs) + + if request_has_perf: + profiled_request_count += 1 + + metrics: dict[str, float] = {} + for key, values in values_by_metric.items(): + if not values: + continue + metrics |= dict_add_prefix(compute_statistics(values), f"{key}/") + + return metrics + + +def _iter_vllm_generate_attrs(all_samples: list[Sample]): + for sample in all_samples: + trace = getattr(sample, "trace", None) + if not isinstance(trace, dict): + continue + for event in trace.get("events") or []: + if event.get("type") != "span_end" or event.get("name") != "vllm_generate": + continue + attrs = event.get("attrs") + if isinstance(attrs, dict): + yield attrs + + +def _compute_zero_std_metrics(args, all_samples: list[Sample]): + # only compute in GRPO-like algorithms where one prompt has multiple responses + if args.advantage_estimator == "ppo": + return {} + + def _is_zero_std(samples: list[Sample]): + rewards = [sample.get_reward_value(args) for sample in samples] + return len(rewards) == 0 or all(rewards[0] == r for r in rewards) + + all_sample_groups = group_by(all_samples, lambda s: s.group_index) + interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)] + + interesting_rewards = [str(round(g[0].get_reward_value(args), 1)) for g in interesting_sample_groups] + + return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} + + +def _compute_top_p_kept_vocab_metrics(all_samples: list[Sample]): + total_kept = 0 + total_tokens = 0 + for sample in all_samples: + offsets = sample.rollout_top_p_token_offsets + if offsets is None or sample.response_length == 0: + continue + offsets = torch.as_tensor(offsets, dtype=torch.int64) + if offsets.numel() == 0: + continue + assert ( + offsets.numel() == sample.response_length + 1 + ), f"top-p token offsets length {offsets.numel()} != response length + 1 {sample.response_length + 1}" + if sample.remove_sample: + continue + if sample.loss_mask is None: + total_kept += int(offsets[-1] - offsets[0]) + total_tokens += sample.response_length + continue + loss_mask = torch.as_tensor(sample.loss_mask, dtype=torch.bool, device=offsets.device) + assert ( + loss_mask.numel() == sample.response_length + ), f"loss mask length {loss_mask.numel()} != response length {sample.response_length}" + total_kept += int(torch.diff(offsets)[loss_mask].sum()) + total_tokens += int(loss_mask.sum()) + if total_tokens == 0: + return {} + return {"top_p_kept_vocab_per_token": total_kept / total_tokens} + + +def _compute_spec_metrics(args, all_samples: list[Sample]): + if getattr(args, "vllm_speculative_algorithm", None) is None: + return {} + num_samples = len(all_samples) + metrics = {} + metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples + metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples + return metrics + + +def _compute_prefix_cache_metrics(all_samples: list[Sample]): + num_samples = len(all_samples) + metrics = {} + total_cached_tokens = sum(sample.prefix_cache_info.cached_tokens for sample in all_samples) + total_prompt_tokens = sum(sample.prefix_cache_info.total_prompt_tokens for sample in all_samples) + + metrics["prefix_cache_hit_rate"] = total_cached_tokens / total_prompt_tokens if total_prompt_tokens > 0 else 0.0 + metrics["avg_cached_tokens_per_sample"] = total_cached_tokens / num_samples + return metrics + + +def _compute_reward_cat_metrics(args, all_samples: list[Sample]): + reward_cat_key = args.log_reward_category + if reward_cat_key is None: + return {} + + samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key]) + + return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()} + + +def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): + if args.custom_eval_rollout_log_function_path is not None: + custom_log_func = load_function(args.custom_eval_rollout_log_function_path) + if custom_log_func(rollout_id, args, data, extra_metrics): + return + + log_dict = extra_metrics or {} + for key in data.keys(): + rewards = data[key]["rewards"] + log_dict[f"eval/{key}"] = sum(rewards) / len(rewards) + if (samples := data[key].get("samples")) is not None: + log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/") + if "truncated" in data[key]: + truncated = data[key]["truncated"] + log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated) + if args.log_passrate: + log_dict |= dict_add_prefix( + compute_pass_rate( + flat_rewards=rewards, + group_size=args.n_samples_per_eval_prompt, + ), + f"eval/{key}-", + ) + + logger.info(f"eval {rollout_id}: {log_dict}") + + step = compute_rollout_step(args, rollout_id) + log_dict["eval/step"] = step + logging_utils.log(args, log_dict, step_key="eval/step") + + return log_dict + + +def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): + if args.custom_rollout_log_function_path is not None: + custom_log_func = load_function(args.custom_rollout_log_function_path) + if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time): + return + + if args.load_debug_rollout_data: + return + + log_dict = {**(rollout_extra_metrics or {})} + log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), "rollout/") + log_dict |= dict_add_prefix(compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/") + logger.info(f"perf {rollout_id}: {log_dict}") + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + logging_utils.log(args, log_dict, step_key="rollout/step") diff --git a/vime/utils/tensorboard_utils.py b/vime/observability/tensorboard_utils.py similarity index 96% rename from vime/utils/tensorboard_utils.py rename to vime/observability/tensorboard_utils.py index af79bb03c..16945dcd8 100644 --- a/vime/utils/tensorboard_utils.py +++ b/vime/observability/tensorboard_utils.py @@ -22,7 +22,7 @@ class _TensorboardAdapter(metaclass=SingletonMeta): # tb.log({"Loss": 0.1}, step=1) # In other files: - # from tensorboard_utils import _TensorboardAdapter + # from vime.observability.tensorboard_utils import _TensorboardAdapter # tb = _TensorboardAdapter(args) # No parameters needed to get existing instance # tb.log({"Accuracy": 0.9}, step=1) """ diff --git a/vime/utils/timer.py b/vime/observability/timer.py similarity index 98% rename from vime/utils/timer.py rename to vime/observability/timer.py index ec1bdf767..d3feb4916 100644 --- a/vime/utils/timer.py +++ b/vime/observability/timer.py @@ -5,7 +5,7 @@ import torch.distributed -from .misc import SingletonMeta +from vime.utils.misc import SingletonMeta __all__ = ["Timer", "timer"] diff --git a/vime/utils/trace_utils.py b/vime/observability/trace_utils.py similarity index 72% rename from vime/utils/trace_utils.py rename to vime/observability/trace_utils.py index 0c0f51aa2..34b4a3e93 100644 --- a/vime/utils/trace_utils.py +++ b/vime/observability/trace_utils.py @@ -14,6 +14,34 @@ from vime.utils.types import Sample TRACE_VERSION = 1 +TRACE_CHILDREN_KEY = "_trace_children" +VLLM_TRACE_META_KEYS = ( + "prompt_tokens", + "completion_tokens", + "cached_tokens", + "queue_time", + "e2e_latency", + "decode_throughput", +) +VLLM_PD_PREFILL_SEGMENTS = ( + ("pd_prefill_bootstrap_queue_duration", "vllm_pd_prefill_bootstrap_queue"), + ("pd_prefill_bootstrap_duration", "vllm_pd_prefill_bootstrap"), + ("pd_prefill_alloc_wait_duration", "vllm_pd_prefill_alloc_wait"), + ("pd_prefill_forward_duration", "vllm_pd_prefill_forward"), + ("pd_prefill_transfer_queue_duration", "vllm_pd_prefill_transfer_queue"), +) +VLLM_PD_DECODE_SEGMENTS = ( + ("pd_decode_prealloc_duration", "vllm_pd_decode_prealloc"), + ("pd_decode_bootstrap_duration", "vllm_pd_decode_bootstrap"), + ("pd_decode_alloc_wait_duration", "vllm_pd_decode_alloc_wait"), + ("pd_decode_transfer_duration", "vllm_pd_decode_transfer"), + ("pd_decode_forward_duration", "vllm_pd_decode_forward"), +) +VLLM_PD_SUMMARY_KEYS = ( + "pd_transfer_speed_gb_s", + "pd_transfer_total_mb", + "pd_prefill_retry_count", +) logger = logging.getLogger(__name__) _TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar( @@ -41,6 +69,8 @@ class TraceHandle: class TraceSpanContext: target: Sample | TraceHandle | list[Sample | TraceHandle] handles: list[TraceHandle] + span_records: list[tuple[TraceHandle, str]] = field(default_factory=list) + start_ts: float = 0.0 end_attrs: dict[str, Any] = field(default_factory=dict) end_events: list[dict[str, Any]] = field(default_factory=list) closed: bool = False @@ -51,17 +81,19 @@ def set(self, key: str, value: Any) -> TraceSpanContext: return self def update(self, attrs: dict[str, Any] | None) -> TraceSpanContext: - if attrs: - self.end_attrs.update(attrs) - self._sync_end_events(attrs) + try: + if attrs: + plain_attrs = dict(attrs) + trace_children = plain_attrs.pop(TRACE_CHILDREN_KEY, None) + if plain_attrs: + self.end_attrs.update(plain_attrs) + self._sync_end_events(plain_attrs) + if trace_children: + _append_trace_children(self.span_records, trace_children, parent_start_ts=self.start_ts) + except Exception as exc: + _log_trace_error("update", exc) return self - def set_attr(self, key: str, value: Any) -> TraceSpanContext: - return self.set(key, value) - - def update_attrs(self, attrs: dict[str, Any] | None) -> TraceSpanContext: - return self.update(attrs) - def build_end_attrs(self) -> dict[str, Any] | None: return dict(self.end_attrs) or None @@ -105,19 +137,79 @@ def _new_span_id() -> str: return uuid.uuid4().hex -def build_vllm_meta_trace_attrs(output: dict[str, Any]) -> dict[str, Any]: - """Trace-span attributes from a vLLM ``/inference/v1/generate`` response.""" +def build_vllm_meta_trace_attrs(meta: dict[str, Any]) -> dict[str, Any]: attrs: dict[str, Any] = {} - choices = output.get("choices") or [] - if choices and choices[0].get("finish_reason") is not None: - attrs["finish_reason"] = choices[0]["finish_reason"] - usage = output.get("usage") or {} - for key in ("prompt_tokens", "completion_tokens", "cached_tokens"): - if usage.get(key) is not None: - attrs[key] = usage[key] + try: + attrs.update({key: meta[key] for key in VLLM_TRACE_META_KEYS if key in meta and meta[key] is not None}) + finish_reason = meta.get("finish_reason") + if isinstance(finish_reason, dict) and finish_reason.get("type") is not None: + attrs["finish_reason"] = finish_reason["type"] + elif finish_reason is not None: + attrs["finish_reason"] = finish_reason + + if meta.get("id") is not None: + attrs["vllm_request_id"] = meta["id"] + + trace_children = _build_vllm_pd_trace_children(meta) + if trace_children: + attrs[TRACE_CHILDREN_KEY] = trace_children + except Exception as exc: + _log_trace_error("vllm_meta_attrs", exc) return attrs +def _build_vllm_pd_trace_children(meta: dict[str, Any]) -> list[dict[str, Any]]: + trace_children: list[dict[str, Any]] = [] + cursor = 0.0 + for phase_name, phase_label, segments in ( + ("vllm_pd_prefill", "prefill", VLLM_PD_PREFILL_SEGMENTS), + ("vllm_pd_decode", "decode", VLLM_PD_DECODE_SEGMENTS), + ): + phase_children: list[dict[str, Any]] = [] + phase_cursor = 0.0 + for key, child_name in segments: + if key not in meta or meta[key] is None: + continue + duration = float(meta[key]) + if duration <= 0.0: + continue + phase_children.append( + { + "type": "span", + "name": child_name, + "start_offset": phase_cursor, + "end_offset": phase_cursor + duration, + "attrs": {key: meta[key]}, + } + ) + phase_cursor += duration + if not phase_children: + continue + trace_children.append( + { + "type": "span", + "name": phase_name, + "start_offset": cursor, + "end_offset": cursor + phase_cursor, + "attrs": {"phase": phase_label, "duration_s": phase_cursor}, + "children": phase_children, + } + ) + cursor += phase_cursor + + summary_attrs = {key: meta[key] for key in VLLM_PD_SUMMARY_KEYS if key in meta and meta[key] is not None} + if summary_attrs: + trace_children.append( + { + "type": "event", + "name": "vllm_pd_summary", + "start_offset": cursor, + "attrs": summary_attrs, + } + ) + return trace_children + + def _ensure_trace_carrier( carrier: dict[str, Any] | None, *, @@ -237,7 +329,14 @@ def trace_event( try: timestamp = time.time() for handle in _coerce_handles(target): - _append_event(handle, kind="event", name=name, timestamp=timestamp, attrs=attrs) + _append_event( + handle, + kind="event", + name=name, + timestamp=timestamp, + parent_span_id=handle.parent_span_id or _get_current_parent_span_id(handle.trace_id), + attrs=attrs, + ) except Exception as exc: _log_trace_error(f"event:{name}", exc) @@ -291,6 +390,8 @@ def trace_span( span_context = TraceSpanContext( target=handles[0] if len(handles) == 1 else handles, handles=handles, + span_records=span_records, + start_ts=timestamp, ) try: @@ -417,6 +518,72 @@ def _record_span_end( return events +def _append_trace_children( + parent_records: list[tuple[TraceHandle, str]], + trace_children: Any, + *, + parent_start_ts: float, +) -> None: + if isinstance(trace_children, dict): + trace_children = [trace_children] + if not isinstance(trace_children, list): + return + + for child in trace_children: + if not isinstance(child, dict): + continue + child_type = child.get("type") + child_name = child.get("name") + if child_type not in ("span", "event") or not child_name: + continue + + child_start_ts = parent_start_ts + float(child["start_offset"]) + attrs = child.get("attrs") if isinstance(child.get("attrs"), dict) else None + + if child_type == "event": + for handle, parent_span_id in parent_records: + _append_event( + handle, + kind="event", + name=child_name, + timestamp=child_start_ts, + parent_span_id=parent_span_id, + attrs=attrs, + ) + continue + + child_end_ts = parent_start_ts + float(child["end_offset"]) + if child_end_ts < child_start_ts: + continue + child_records: list[tuple[TraceHandle, str]] = [] + for handle, parent_span_id in parent_records: + child_span_id = _new_span_id() + _append_event( + handle, + kind="span_start", + name=child_name, + timestamp=child_start_ts, + span_id=child_span_id, + parent_span_id=parent_span_id, + ) + child_records.append((handle, child_span_id)) + + _append_trace_children( + child_records, + child.get("children"), + parent_start_ts=child_start_ts, + ) + for handle, child_span_id in child_records: + _append_event( + handle, + kind="span_end", + name=child_name, + timestamp=child_end_ts, + span_id=child_span_id, + attrs=attrs, + ) + + def _append_event( handle: TraceHandle, *, diff --git a/vime/observability/train_data_utils.py b/vime/observability/train_data_utils.py new file mode 100644 index 000000000..54b9b7937 --- /dev/null +++ b/vime/observability/train_data_utils.py @@ -0,0 +1,242 @@ +import logging +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + + +_CONTEXT_PARALLEL_FIELDS = ( + "rollout_log_probs", + "teacher_log_probs", + "log_probs", + "ref_log_probs", + "values", + "advantages", + "returns", + "kl", + "entropy", + "opd_reverse_kl", +) + + +def _to_cpu(value): + if torch.is_tensor(value): + return value.detach().cpu() + if isinstance(value, dict): + return {key: _to_cpu(item) for key, item in value.items()} + if isinstance(value, list): + return [_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(_to_cpu(item) for item in value) + return value + + +@torch.no_grad() +def restore_context_parallel_fields_to_cpu( + rollout_data: dict[str, Any], + gather_tensor: Callable[[torch.Tensor, int, int], torch.Tensor], + *, + keep_restored: bool, +) -> dict[str, Any] | None: + """Restore CP fields one tensor at a time, retaining CPU values only on the writer.""" + total_lengths = rollout_data["total_lengths"] + response_lengths = rollout_data["response_lengths"] + num_samples = len(response_lengths) + if len(total_lengths) != num_samples: + raise ValueError( + "total_lengths and response_lengths must contain the same number of samples, " + f"got {len(total_lengths)} and {num_samples}." + ) + + restored = ( + {key: _to_cpu(value) for key, value in rollout_data.items() if key not in _CONTEXT_PARALLEL_FIELDS} + if keep_restored + else None + ) + for key in _CONTEXT_PARALLEL_FIELDS: + values = rollout_data.get(key) + if values is None: + continue + if not isinstance(values, (list, tuple)) or len(values) != num_samples: + value_count = len(values) if isinstance(values, (list, tuple)) else "not a list or tuple" + raise ValueError( + f"CP-sharded field {key!r} must contain one tensor per sample, " + f"got {value_count}; expected {num_samples}." + ) + + full_values = [] if keep_restored else None + for sample_id, (value, total_length, response_length) in enumerate( + zip(values, total_lengths, response_lengths, strict=True) + ): + if not torch.is_tensor(value): + raise TypeError( + f"CP-sharded field {key!r} sample {sample_id} must be a tensor, got {type(value).__name__}." + ) + full_value = gather_tensor(value, int(total_length), int(response_length)).detach() + if full_value.size(0) != response_length: + raise ValueError( + f"Restored field {key!r} sample {sample_id} has length {full_value.size(0)}, " + f"expected {response_length}." + ) + if keep_restored: + assert full_values is not None + full_values.append(full_value.detach().cpu()) + # Drop the full GPU tensor before the next collective allocates its + # output, keeping peak extra device memory to one response tensor. + del full_value + if keep_restored: + assert restored is not None + restored[key] = full_values + + return restored + + +# Per-rank scheduling arrangement (not per-sample); stored under `dp_shards` so +# the flat `samples` view can stay aligned with the rollout debug dump. +_LAYOUT_FIELDS = ("micro_batch_indices", "num_microbatches", "global_batch_sizes") +# Whole-batch fields that are identical on every DP shard; stored once at the top. +_WHOLE_BATCH_FIELDS = ("raw_reward",) + + +def _is_per_sample(value, num_samples: int) -> bool: + if torch.is_tensor(value): + return value.dim() >= 1 and value.size(0) == num_samples + if isinstance(value, (list, tuple)): + return len(value) == num_samples + return False + + +def _build_dump_payload(dp_shards, *, rollout_id, writer_rank): + """Turn the gathered per-rank shards into a rollout-dump-aligned payload. + + Mirrors the rollout debug dump: a ``samples`` list (one dict per training + sample, sorted so it lines up with the rollout dump's ``samples``), plus a + parallel ``dp_shards`` key that keeps the DP/mbs layout (``sample_indices``, + ``partition`` + ``micro_batch_indices`` / ``num_microbatches`` / + ``global_batch_sizes``) without duplicating any per-sample tensor. + + Ordering key preference: ``partition`` (each sample's index in the rollout + debug dump, always available when dumping) restores exact rollout order + regardless of ``sample.index``; ``sample_indices`` (``sample.index``, may be + ``None``) is the fallback; failing both, DP-gather order is kept. + """ + samples = [] + layout = [] + whole_batch = {} + # Per-sample id columns handled specially: pulled out of the generic + # transpose and surfaced as scalar keys on each sample dict. + id_columns = {"sample_indices": "sample_index", "partition": "rollout_position"} + for shard in dp_shards: + rollout_data = shard["rollout_data"] + dp_rank = shard["data_parallel_rank"] + num_local = len(rollout_data["response_lengths"]) + + shard_layout = {"rank": shard["rank"], "data_parallel_rank": dp_rank} + for column in id_columns: + value = rollout_data.get(column) + shard_layout[column] = list(value) if value is not None else None + per_sample_columns = {} + for key, value in rollout_data.items(): + if key in id_columns: + continue + if key in _WHOLE_BATCH_FIELDS: + whole_batch.setdefault(key, value) + elif key in _LAYOUT_FIELDS or not _is_per_sample(value, num_local): + # Scheduling / non-per-sample fields stay in the layout so the + # flat sample view holds only what pairs with a single sample. + shard_layout[key] = value + else: + per_sample_columns[key] = value + + for i in range(num_local): + sample = {"data_parallel_rank": dp_rank} + for column, scalar_key in id_columns.items(): + values = rollout_data.get(column) + if values is not None: + sample[scalar_key] = values[i] + for key, column in per_sample_columns.items(): + sample[key] = column[i] + samples.append(sample) + layout.append(shard_layout) + + order_key = next( + ( + scalar_key + for scalar_key in ("rollout_position", "sample_index") + if samples and all(sample.get(scalar_key) is not None for sample in samples) + ), + None, + ) + if order_key is not None: + samples.sort(key=lambda sample: sample[order_key]) + else: + logger.warning( + "Cannot restore global sample order for the train debug dump: samples carry neither " + "partition nor sample_index. Saving samples in DP-gather order instead." + ) + + return { + "format_version": 2, + "rollout_id": rollout_id, + "rank": writer_rank, + "samples": samples, + "dp_shards": layout, + **whole_batch, + } + + +def save_debug_train_data(args, *, rollout_id, rollout_data): + path_template = args.save_debug_train_data + if path_template is None: + return + + from megatron.core import mpu + + # The last PP stage owns the computed train fields. TP ranks hold the same + # token values after TP reduction, so only TP0 needs to participate below. + if not mpu.is_pipeline_last_stage(ignore_virtual=True) or mpu.get_tensor_model_parallel_rank() != 0: + return + + from vime.backends.megatron_utils.cp_utils import all_gather_with_cp + + # All CP ranks in the selected PP/TP group must enter these collectives. + # CP=1 follows the same normalization path but all_gather_with_cp is an identity. + cp_rank = mpu.get_context_parallel_rank() + rollout_data = restore_context_parallel_fields_to_cpu( + rollout_data, + all_gather_with_cp, + keep_restored=cp_rank == 0, + ) + if cp_rank != 0: + return + assert rollout_data is not None + + rank = torch.distributed.get_rank() + local_shard = { + "rank": rank, + "data_parallel_rank": mpu.get_data_parallel_rank(with_context_parallel=False), + "rollout_data": rollout_data, + } + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + dp_src_rank = mpu.get_data_parallel_src_rank(with_context_parallel=False) + if dp_size == 1: + dp_shards = [local_shard] + else: + dp_shards = [None] * dp_size if rank == dp_src_rank else None + torch.distributed.gather_object( + local_shard, + dp_shards, + dst=dp_src_rank, + group=mpu.get_data_parallel_group_gloo(with_context_parallel=False), + ) + + if rank != dp_src_rank: + return + + path = Path(path_template.format(rollout_id=rollout_id, rank=rank)) + logger.info(f"Save debug train data from {dp_size} DP shard(s) to {path}") + path.parent.mkdir(parents=True, exist_ok=True) + torch.save(_build_dump_payload(dp_shards, rollout_id=rollout_id, writer_rank=rank), path) diff --git a/vime/observability/train_metric_utils.py b/vime/observability/train_metric_utils.py new file mode 100644 index 000000000..0d42c1bcb --- /dev/null +++ b/vime/observability/train_metric_utils.py @@ -0,0 +1,405 @@ +import logging +from argparse import Namespace +from copy import deepcopy + +import numpy as np +import torch +import torch.distributed as dist + +from vime.observability import logging_utils +from vime.observability.metric_utils import compute_pass_rate, compute_rollout_step +from vime.observability.timer import Timer +from vime.utils.flops_utils import calculate_fwd_flops +from vime.utils.types import RolloutBatch + +logger = logging.getLogger(__name__) + + +def reduce_train_step_metrics( + losses_reduced: list[dict], + *, + calculate_per_token_loss: bool, + step_global_batch_size: int, + cp_size: int, + dp_with_cp_group, +) -> dict[str, float]: + """Aggregate per-mb log dicts into the dict ``train_one_step`` reports. + + Pipeline (1:1 with what the train loop used to do inline): + 1. Sum each metric's per-mb ``values`` tensor locally on this rank. + 2. All-reduce across the DP*CP group (``dp_with_cp_group``). + 3. Apply the per-mode divisor / cp_factor: + - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``, + CP-inflated by ``cp_size`` because every CP rank computes the same + num_tokens off the FULL (not chunked) masks; the + ``cp_factor = cp_size`` multiplier cancels that inflation, leaving + the genuine per-token average. + - per-rollout-mean: divisor = constant ``step_global_batch_size`` from + the rollout side, never all-reduced, so no CP inflation to cancel + and ``cp_factor = 1``. + + Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce`` + to a no-op, then pre-aggregate virtual ranks themselves — this exercises + the same call shape as production while staying single-process. + """ + keys = losses_reduced[0]["keys"] + values = None + for item in losses_reduced: + values = item["values"] if values is None else values + item["values"] + assert len(keys) + 1 == values.numel() + dist.all_reduce(values, group=dp_with_cp_group) + values = values.tolist() + + if calculate_per_token_loss: + num_samples_or_tokens = values[0] + cp_factor = cp_size + else: + num_samples_or_tokens = step_global_batch_size + cp_factor = 1 + return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)} + + +def rollout_log_metric_contribution( + per_rank_reducer_sum: float, + *, + cp_size: int, + num_rollouts_in_rollout: int, + dp_size: int, +) -> tuple[float, float]: + """``(sum, count)`` tuple for a per-rollout-mean metric. + + Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout`` + (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size * + cp_size`` ranks, and each rank emits the same ``count``, so the totals + cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount = + sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports + for the same samples (when ``num_steps_per_rollout == 1``). + + Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end + in tests. + """ + sum_value = cp_size * per_rank_reducer_sum + count = num_rollouts_in_rollout / dp_size + return sum_value, count + + +def gather_and_reduce_log_dict( + log_dict: dict, + *, + dp_size: int, + dp_src_rank: int, + dp_group, +) -> dict | None: + """Gather per-rank log dicts and reduce each metric on ``dp_src_rank``. + + ``(sum, count)`` tuples reduce to ``Σsum / Σcount``; plain values reduce + to a mean across ranks. The helper stays free of reporting side effects so + CPU multi-process tests can exercise it with real ``torch.distributed``. + """ + if dist.get_rank() == dp_src_rank: + gathered = [None] * dp_size + dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group) + reduced: dict = {} + for key in log_dict: + values = [item[key] for item in gathered] + first = values[0] + if isinstance(first, tuple) and len(first) == 2: + total_sum = sum(value[0] for value in values) + total_count = sum(value[1] for value in values) + reduced[key] = total_sum / total_count if total_count else 0.0 + else: + reduced[key] = sum(values) / dp_size + return reduced + dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group) + return None + + +def gather_log_data( + metric_name: str, + args: Namespace, + rollout_id: int, + log_dict: dict[str, "float | tuple[float, float]"], +) -> dict[str, float] | None: + """Gather per-rank metrics and report them through the configured trackers.""" + from megatron.core import mpu + + reduced = gather_and_reduce_log_dict( + log_dict, + dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True), + dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True), + dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), + ) + if reduced is None: + return None + reduced_log_dict = {f"{metric_name}/{key}": value for key, value in reduced.items()} + logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") + step = compute_rollout_step(args, rollout_id) + reduced_log_dict["rollout/step"] = step + logging_utils.log(args, reduced_log_dict, step_key="rollout/step") + return reduced_log_dict + + +def log_rollout_data( + rollout_id: int, + args: Namespace, + rollout_data: RolloutBatch, +) -> None: + """Summarize and report Megatron-side rollout fields.""" + from megatron.core import mpu + + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean + + if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): + cp_size = mpu.get_context_parallel_world_size() + log_dict = {} + response_lengths = rollout_data["response_lengths"] + loss_masks = rollout_data["loss_masks"] + total_lengths = rollout_data["total_lengths"] + rollout_mask_sums = rollout_data.get("rollout_mask_sums", None) + dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False) + num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"]) + + ignored_keys = { + "tokens", + "multimodal_train_inputs", + "loss_masks", + "sample_indices", + "rollout_ids", + "rollout_mask_sums", + "rollout_top_p_token_ids", + "rollout_top_p_token_offsets", + "rollout_routed_experts", + "global_batch_sizes", + "num_microbatches", + "micro_batch_indices", + "source_names", + "local_raw_reward", + } + per_rollout_mean_keys = { + "log_probs", + "ref_log_probs", + "rollout_log_probs", + "returns", + "advantages", + "values", + "teacher_log_probs", + "opd_reverse_kl", + } + + for key, value in rollout_data.items(): + if key in ignored_keys: + continue + if isinstance(value, (list, tuple)): + count = len(value) + if isinstance(value[0], torch.Tensor): + tensor = torch.cat(value).clone().detach() + if key in per_rollout_mean_keys: + sum_of_sample_mean = get_sum_of_sample_mean( + total_lengths, + response_lengths, + loss_masks, + rollout_mask_sums, + ) + sum_value, count = rollout_log_metric_contribution( + sum_of_sample_mean(tensor).item(), + cp_size=cp_size, + num_rollouts_in_rollout=num_rollouts_in_rollout, + dp_size=dp_world, + ) + log_dict[key] = (sum_value, count) + continue + per_rank_sum = tensor.mean() * cp_size * count + sum_value = per_rank_sum.item() + else: + sum_value = sum(value) + log_dict[key] = (sum_value, count) + elif isinstance(value, torch.Tensor): + log_dict[key] = (value.float().mean().item(), 1) + else: + raise ValueError(f"Unsupported type: {type(value)} for key: {key}") + + reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict) + if args.ci_test and reduced_log_dict is not None: + if ( + rollout_id == 0 + and not getattr(args, "ci_disable_kl_checker", False) + and not getattr(args, "use_rollout_routing_replay", False) + and "rollout/log_probs" in reduced_log_dict + and "rollout/ref_log_probs" in reduced_log_dict + ): + assert abs(reduced_log_dict["rollout/log_probs"] - reduced_log_dict["rollout/ref_log_probs"]) < 1e-8 + if "rollout/log_probs" in reduced_log_dict: + assert -1 < reduced_log_dict["rollout/log_probs"] < 0 + if "rollout/entropy" in reduced_log_dict: + assert 0 < reduced_log_dict["rollout/entropy"] < 1 + + if args.log_multi_turn: + log_multi_turn_data(rollout_id, args, rollout_data) + if args.log_passrate: + log_passrate(rollout_id, args, rollout_data) + + if args.log_correct_samples and mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): + response_lengths = rollout_data["response_lengths"] + loss_masks = rollout_data["loss_masks"] + total_lengths = rollout_data["total_lengths"] + + def quantile(total_value, n_quantiles, data) -> dict: + import math + + assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1." + quantiles = [(i + 1) / n_quantiles for i in range(n_quantiles)] + cut_points = [total_value * quantile for quantile in quantiles] + cut_points[-1] = total_value + + count = [0] * n_quantiles + for value in data: + for i, point in enumerate(cut_points): + if value <= point: + count[i] += 1 + break + + total = sum(count) + 1e-9 + percentile = [value / total for value in count] + return { + f"p{min(math.ceil(quantile * 100), 100)}": value + for quantile, value in zip(quantiles, percentile, strict=True) + } + + raw_rewards = rollout_data["local_raw_reward"] + correct_response_lengths = [] + correct_total_lengths = [] + correct_loss_masks = [] + correct_entropy = [] + for i, raw_reward in enumerate(raw_rewards): + if raw_reward == 1: + correct_response_lengths.append(response_lengths[i]) + correct_total_lengths.append(total_lengths[i]) + correct_loss_masks.append(loss_masks[i]) + correct_entropy.append(-rollout_data["log_probs"][i]) + num_correct_responses = len(correct_total_lengths) + rollout_data["correct_response_lengths"] = correct_response_lengths + correct_response_length_percentile = quantile( + args.rollout_max_response_len, + 4, + rollout_data["correct_response_lengths"], + ) + for percentile, value in correct_response_length_percentile.items(): + rollout_data[f"correct_length/{percentile}"] = [value] * num_correct_responses + if correct_entropy: + sum_of_sample_mean = get_sum_of_sample_mean( + correct_total_lengths, + correct_response_lengths, + correct_loss_masks, + sample_denoms=None, + ) + correct_entropy_value = sum_of_sample_mean(torch.cat(correct_entropy, dim=0)) + rollout_data["correct_entropy"] = [correct_entropy_value.item()] * num_correct_responses + else: + rollout_data["correct_entropy"] = [0] * num_correct_responses + + +def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None: + """Report multi-turn response-length and round-count metrics.""" + from megatron.core import mpu + + if mpu.get_tensor_model_parallel_rank() != 0 or not mpu.is_pipeline_last_stage(): + return + + log_dict = {} + for key, value in rollout_data.items(): + if key == "loss_masks" and value: + device = value[0].device + raw_response_lengths = torch.tensor( + [item.shape[0] for item in value], + dtype=torch.float32, + device=device, + ) + log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item() + log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item() + log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item() + log_dict["raw_response_length/response_length_clip_ratio"] = ( + (raw_response_lengths >= args.rollout_max_response_len).float().mean().item() + ) + + wo_obs_response_lengths = torch.tensor( + [item.sum().item() for item in value], + dtype=torch.float32, + device=device, + ) + log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item() + log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item() + log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item() + if key == "round_number": + round_number_array = np.array(value) + log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array) + log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array) + log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array) + gather_log_data("multi_turn", args, rollout_id, log_dict) + + +def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None: + """Compute and report pass@k metrics from grouped ``raw_reward`` values.""" + from megatron.core import mpu + + if mpu.get_tensor_model_parallel_rank() != 0 or not mpu.is_pipeline_last_stage(): + return + + log_dict = {} + for key, value in rollout_data.items(): + if key == "raw_reward": + log_dict |= compute_pass_rate( + flat_rewards=value, + group_size=args.n_samples_per_prompt, + num_groups=args.rollout_batch_size, + ) + gather_log_data("passrate", args, rollout_id, log_dict) + + +def log_perf_data( + rollout_id: int, + args: Namespace, + extra_metrics: dict | None = None, +) -> None: + from megatron.core import mpu + + timer_instance = Timer() + log_dict_raw = deepcopy(timer_instance.log_dict()) + timer_instance.reset() + + if not ( + mpu.get_tensor_model_parallel_rank() == 0 + and mpu.is_pipeline_last_stage() + and mpu.get_data_parallel_rank(with_context_parallel=True) == 0 + ): + return + + log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()} + if extra_metrics: + log_dict.update(extra_metrics) + + if "perf/actor_train_time" in log_dict: + total_fwd_flops = ( + calculate_fwd_flops(seqlens=timer_instance.seq_lens, args=args) / dist.get_world_size() / 1e12 + ) + + if "perf/log_probs_time" in log_dict: + log_dict["perf/log_probs_tflops"] = total_fwd_flops / log_dict["perf/log_probs_time"] + + if "perf/ref_log_probs_time" in log_dict: + log_dict["perf/ref_log_probs_tflops"] = total_fwd_flops / log_dict["perf/ref_log_probs_time"] + + if log_dict["perf/actor_train_time"] > 0: + log_dict["perf/actor_train_tflops"] = 3 * total_fwd_flops / log_dict["perf/actor_train_time"] + log_dict["perf/actor_train_tok_per_s"] = sum(timer_instance.seq_lens) / log_dict["perf/actor_train_time"] + + if "perf/train_wait_time" in log_dict and "perf/train_time" in log_dict: + total_time = log_dict["perf/train_wait_time"] + log_dict["perf/train_time"] + if total_time > 0: + log_dict["perf/step_time"] = total_time + log_dict["perf/wait_time_ratio"] = log_dict["perf/train_wait_time"] / total_time + + logger.info(f"perf {rollout_id}: {log_dict}") + + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + logging_utils.log(args, log_dict, step_key="rollout/step") diff --git a/vime/utils/wandb_utils.py b/vime/observability/wandb_utils.py similarity index 77% rename from vime/utils/wandb_utils.py rename to vime/observability/wandb_utils.py index 6e4410d1f..cda03013e 100644 --- a/vime/utils/wandb_utils.py +++ b/vime/observability/wandb_utils.py @@ -79,56 +79,6 @@ def init_wandb_primary(args): args.wandb_run_id = wandb.run.id -def reinit_wandb_primary_with_open_metrics(args, router_addr): - """Re-initialize the primary W&B run with open metrics endpoints. - - The primary wandb init happens before rollout servers start (to obtain - ``wandb_run_id`` for secondary processes). This function is called - *after* servers are up so the router address is available for scraping - vLLM Prometheus metrics via the primary process's stats monitor. - """ - if not args.use_wandb or _is_offline_mode(args): - return - if getattr(args, "wandb_mode", None) == "disabled": - return - if router_addr is None: - return - wandb_run_id = getattr(args, "wandb_run_id", None) - if wandb_run_id is None: - return - - logger.info(f"Re-initializing primary W&B with vLLM metrics at {router_addr}.") - - wandb.finish() - - init_kwargs = { - "id": wandb_run_id, - "entity": args.wandb_team, - "project": args.wandb_project, - "resume": "allow", - "reinit": True, - "settings": wandb.Settings( - mode="shared", - x_primary=True, - x_stats_open_metrics_endpoints={ - # router_addr already includes the /metrics path on the vllm-router - # prometheus port (see RolloutManager._get_metrics_router_addr). - "vllm_engine": router_addr, - }, - x_stats_open_metrics_filters={ - "vllm_engine.*": {}, - }, - ), - } - - if args.wandb_dir: - os.makedirs(args.wandb_dir, exist_ok=True) - init_kwargs["dir"] = args.wandb_dir - - wandb.init(**init_kwargs) - _init_wandb_common() - - def _compute_config_for_logging(args): output = _args_to_config_dict(args) diff --git a/vime/platforms/__init__.py b/vime/platforms/__init__.py new file mode 100644 index 000000000..8cec41787 --- /dev/null +++ b/vime/platforms/__init__.py @@ -0,0 +1,83 @@ +"""Accelerator platform discovery and narrow capability providers. + +``VIME_PLATFORM`` is the explicit override. When it is not set, NPU detection +is lazy and failure-safe; CUDA is the default. +""" + +from __future__ import annotations + +import os +from functools import cache + +from .base import ( + CheckpointCapabilities, + Platform, + RayResourceSpec, + TrainingBootstrap, + VLLMLaunchPlatformOps, + WeightTransferPlatformOps, +) +from .cuda import create_cuda_platform +from .npu import create_npu_platform, detect_npu + +_PLATFORM_FACTORIES = { + "cuda": create_cuda_platform, + "npu": create_npu_platform, +} + + +@cache +def get_platform(name: str) -> Platform: + normalized = name.strip().lower() + try: + factory = _PLATFORM_FACTORIES[normalized] + except KeyError as exc: + available = ", ".join(_PLATFORM_FACTORIES) + raise ValueError(f"Unknown Vime platform {name!r}; registered platforms: {available}") from exc + return factory() + + +@cache +def _resolve_platform(override: str | None) -> Platform: + if override: + return get_platform(override) + + try: + if detect_npu(): + return get_platform("npu") + except Exception: # noqa: BLE001 - a failed detector must not break imports + pass + return get_platform("cuda") + + +def current_platform() -> Platform: + """Resolve the active platform without probing hardware at module import.""" + raw_override = os.environ.get("VIME_PLATFORM") + override = raw_override.strip().lower() if raw_override and raw_override.strip() else None + accelerator_override = os.environ.get("VIME_ACCELERATOR", "").strip().lower() + if accelerator_override in {"npu", "cuda", "musa"}: + accelerator_platform = "npu" if accelerator_override == "npu" else "cuda" + if override in _PLATFORM_FACTORIES and override != accelerator_platform: + raise ValueError(f"Conflicting VIME_PLATFORM={override!r} and VIME_ACCELERATOR={accelerator_override!r}") + if override is None: + override = accelerator_platform + return _resolve_platform(override) + + +def reset_platform_cache() -> None: + """Clear resolver/factory caches (primarily for tests and plugin registration).""" + get_platform.cache_clear() + _resolve_platform.cache_clear() + + +__all__ = [ + "CheckpointCapabilities", + "Platform", + "RayResourceSpec", + "TrainingBootstrap", + "VLLMLaunchPlatformOps", + "WeightTransferPlatformOps", + "current_platform", + "get_platform", + "reset_platform_cache", +] diff --git a/vime/platforms/base.py b/vime/platforms/base.py new file mode 100644 index 000000000..c0775e5db --- /dev/null +++ b/vime/platforms/base.py @@ -0,0 +1,143 @@ +"""Small contracts for behavior that genuinely differs by platform.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class RayResourceSpec: + """How one accelerator is represented and addressed by Ray.""" + + resource_name: str + visible_devices_env: str + uses_ray_gpu_resource: bool + + def bundle_resources(self, device_count: float = 1, cpu_count: float = 1) -> dict[str, float]: + return {self.resource_name: device_count, "CPU": cpu_count} + + def actor_options(self, fraction: float) -> dict[str, object]: + if self.uses_ray_gpu_resource: + return {"num_gpus": fraction} + return {"resources": {self.resource_name: fraction}} + + def accelerator_ids(self) -> list[str]: + import ray + + if self.uses_ray_gpu_resource: + ids = ray.get_gpu_ids() + else: + ids = ray.get_runtime_context().get_accelerator_ids().get(self.resource_name, []) + return [str(device_id) for device_id in ids] + + def local_device_id(self) -> int | str: + device_ids = self.accelerator_ids() + if not device_ids: + raise RuntimeError(f"No {self.resource_name} accelerator IDs are assigned to this Ray actor") + + assigned_id = device_ids[0] + visible_devices = os.environ.get(self.visible_devices_env) + if visible_devices is None: + try: + return int(assigned_id) + except ValueError: + return assigned_id + + visible_ids = [value.strip() for value in visible_devices.split(",") if value.strip()] + try: + return visible_ids.index(assigned_id) + except ValueError as exc: + raise RuntimeError( + f"Ray assigned {self.resource_name} id {assigned_id}, but it is absent from " + f"{self.visible_devices_env}={visible_devices!r}" + ) from exc + + def train_runtime_env( + self, + args: Any, + env_vars: Mapping[str, str] | None = None, + ) -> dict[str, str]: + return dict(env_vars or {}) + + def rollout_runtime_env( + self, + args: Any, + env_vars: Mapping[str, str] | None = None, + ) -> dict[str, str]: + return dict(env_vars or {}) + + +class WeightTransferPlatformOps: + """Select vendor backends without changing the shared trainer lifecycle.""" + + def backend(self, backend: str) -> str: + return backend + + def trainer_init_info(self, *, colocate: bool, **kwargs): + if colocate: + from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo + + return IPCTrainerInitInfo(**kwargs) + from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo + + return NCCLTrainerInitInfo(**kwargs) + + +class VLLMLaunchPlatformOps: + """Platform additions to the common vLLM launch command and environment.""" + + def subprocess_env( + self, + base_env: Mapping[str, str], + *, + visible_devices: str, + colocate: bool, + ) -> dict[str, str]: + return dict(base_env) + + +class TrainingBootstrap: + """Lazy Megatron/vendor initialization hooks.""" + + def bootstrap(self) -> None: + return None + + def repatch(self, args: Any) -> None: + return None + + def adjust_tp_partition_dim(self, name: str, partition_dim: int) -> int: + return partition_dim + + def training_context(self, offload_train: bool): + return nullcontext() + + def initialize_optimizer_state(self, optimizer: Any) -> None: + return None + + +@dataclass(frozen=True) +class CheckpointCapabilities: + default_megatron_to_hf_mode: str = "raw" + + def patch_default_planner(self, default_planner: Any) -> None: + return None + + +@dataclass(frozen=True) +class Platform: + """Aggregate only the providers whose semantics differ on Ascend.""" + + name: str + ray: RayResourceSpec + weight_transfer: WeightTransferPlatformOps + vllm: VLLMLaunchPlatformOps + megatron: TrainingBootstrap + checkpoint: CheckpointCapabilities + + @property + def is_npu(self) -> bool: + return self.name == "npu" diff --git a/vime/platforms/cuda.py b/vime/platforms/cuda.py new file mode 100644 index 000000000..b068eebab --- /dev/null +++ b/vime/platforms/cuda.py @@ -0,0 +1,25 @@ +"""Default CUDA platform assembled from the shared CUDA-compatible behavior.""" + +from __future__ import annotations + +from .base import ( + CheckpointCapabilities, + Platform, + RayResourceSpec, + TrainingBootstrap, + VLLMLaunchPlatformOps, + WeightTransferPlatformOps, +) + + +def create_cuda_platform() -> Platform: + return Platform( + name="cuda", + ray=RayResourceSpec( + resource_name="GPU", visible_devices_env="CUDA_VISIBLE_DEVICES", uses_ray_gpu_resource=True + ), + weight_transfer=WeightTransferPlatformOps(), + vllm=VLLMLaunchPlatformOps(), + megatron=TrainingBootstrap(), + checkpoint=CheckpointCapabilities(), + ) diff --git a/vime/platforms/npu.py b/vime/platforms/npu.py new file mode 100644 index 000000000..94d2fea2a --- /dev/null +++ b/vime/platforms/npu.py @@ -0,0 +1,292 @@ +"""Ascend NPU implementation of the Vime platform contracts.""" + +from __future__ import annotations + +import importlib +import logging +import os +from contextlib import nullcontext +from glob import glob +from typing import Any + +from vime.utils import accelerator +from vime.utils.accelerator.torch_accelerator import TorchAccelerator + +from .base import ( + CheckpointCapabilities, + Platform, + RayResourceSpec, + TrainingBootstrap, + VLLMLaunchPlatformOps, + WeightTransferPlatformOps, +) + +logger = logging.getLogger(__name__) + + +class NPUAccelerator(TorchAccelerator): + name = "npu" + device_type = "npu" + communication_backend_name = "hccl" + + def _module(self): + return getattr(importlib.import_module("torch"), "npu", None) + + @property + def visible_devices_env(self) -> str: + return "ASCEND_RT_VISIBLE_DEVICES" + + def distributed_device_id(self, index=None): + # Preserve lazy HCCL initialization instead of CUDA's eager device binding. + return None + + def set_allocator_expandable_segments(self) -> bool: + # NPU allocator policy belongs to the existing TMS/runtime-env hooks. + return False + + +def register_npu_accelerator() -> None: + accelerator.register_accelerator( + "npu", + NPUAccelerator, + is_available=lambda: os.environ.get("VIME_PLATFORM", "").strip().lower() != "cuda" + and NPUAccelerator().is_available(), + priority=300, + communication_backends=("hccl",), + ) + + +def detect_npu() -> bool: + """Return whether a usable NPU is visible, without leaking probe errors.""" + # Do not import torch/torch_npu on an unselected CUDA host merely because + # torch_npu happens to be installed. Its import has process-wide monkey + # patch side effects. An explicit VIME_PLATFORM=npu override bypasses + # detection, while automatic selection first requires an exposed device. + if not (os.path.exists("/dev/davinci_manager") or glob("/dev/davinci[0-9]*")): + return False + try: + torch = importlib.import_module("torch") + if getattr(torch, "npu", None) is None: + importlib.import_module("torch_npu") + npu = getattr(torch, "npu", None) + return bool(npu is not None and npu.is_available()) + except Exception: # noqa: BLE001 - detection must be safe on non-NPU hosts + return False + + +def _ensure_torch_npu() -> None: + importlib.import_module("torch_npu") + + +def _install_safe_empty_cache() -> None: + """Preserve the Ascend allocator guard required by MindSpeed/TMS callers.""" + torch = importlib.import_module("torch") + original_empty_cache = torch.npu.empty_cache + if not getattr(original_empty_cache, "_vime_safe_empty_cache", False): + + def _safe_empty_cache(_original=original_empty_cache) -> None: + try: + _original() + except RuntimeError: + pass + + _safe_empty_cache._vime_safe_empty_cache = True + torch.npu.empty_cache = _safe_empty_cache + + # Some shared dependencies still call the CUDA spelling after torch_npu's + # compatibility patching. Keep that alias local to NPU-bootstrapped jobs. + torch.cuda.empty_cache = torch.npu.empty_cache + + +def _cann_python_site_packages() -> str | None: + candidates: list[str] = [] + for env_key in ("ASCEND_TOOLKIT_HOME", "ASCEND_HOME_PATH"): + base = os.environ.get(env_key) + if not base: + continue + candidates.extend( + [ + os.path.join(base, "python", "site-packages"), + os.path.normpath(os.path.join(base, "..", "python", "site-packages")), + ] + ) + candidates.append("/usr/local/Ascend/ascend-toolkit/latest/python/site-packages") + for path in candidates: + if os.path.isdir(os.path.join(path, "acl")): + return path + return None + + +def _prepend_pythonpath(env: dict[str, str], *paths: str) -> None: + existing = env.get("PYTHONPATH", os.environ.get("PYTHONPATH", "")) + existing_parts = {part for part in existing.split(os.pathsep) if part} + prefix_parts = [path for path in paths if path and path not in existing_parts] + if prefix_parts: + env["PYTHONPATH"] = os.pathsep.join([*prefix_parts, existing] if existing else prefix_parts) + + +class NpuRayResourceSpec(RayResourceSpec): + def __init__(self) -> None: + super().__init__( + resource_name="NPU", + visible_devices_env="ASCEND_RT_VISIBLE_DEVICES", + uses_ray_gpu_resource=False, + ) + + def train_runtime_env(self, args: Any, env_vars=None) -> dict[str, str]: + env = dict(env_vars or {}) + if not (getattr(args, "offload_train", False) and getattr(args, "train_backend", None) == "megatron"): + return env + + env["TMS_HOOK_MODE"] = "torch" + env["TMS_REGION_TAG"] = "training" + env["TMS_ENABLE_CPU_BACKUP"] = "1" + if getattr(args, "colocate", False): + env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" + cann_python_path = _cann_python_site_packages() + if cann_python_path is not None: + _prepend_pythonpath(env, cann_python_path) + return env + + def rollout_runtime_env(self, args: Any, env_vars=None) -> dict[str, str]: + env = dict(env_vars or {}) + cann_python_path = _cann_python_site_packages() + if cann_python_path is not None: + _prepend_pythonpath(env, cann_python_path) + if getattr(args, "colocate", False): + env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" + env["VLLM_USE_AOT_COMPILE"] = "0" + return env + + +class NpuWeightTransferPlatformOps(WeightTransferPlatformOps): + def current_device_uuid(self) -> str: + # Reuse vLLM Ascend's canonical host-IP/physical-chip identifier so + # trainer and receiver always use the exact same mapping. + from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import npu_generate_uuid + + return npu_generate_uuid() + + def backend(self, backend: str) -> str: + return {"ipc": "npu_ipc", "nccl": "hccl"}.get(backend, backend) + + def trainer_init_info(self, *, colocate: bool, **kwargs): + _ensure_torch_npu() + from vllm.plugins import load_general_plugins + + # Trainers, unlike vLLM workers, may not have loaded general plugins yet. + load_general_plugins() + if colocate: + from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import NPUIPCTrainerInitInfo + + return NPUIPCTrainerInitInfo(**kwargs) + from vllm_ascend.distributed.weight_transfer.hccl_engine import HCCLTrainerInitInfo + + return HCCLTrainerInitInfo(**kwargs) + + +class NpuVLLMLaunchPlatformOps(VLLMLaunchPlatformOps): + def subprocess_env(self, base_env, *, visible_devices: str, colocate: bool) -> dict[str, str]: + env = dict(base_env) + env.pop("PYTORCH_CUDA_ALLOC_CONF", None) + env.pop("CUDA_VISIBLE_DEVICES", None) + env.pop("HIP_VISIBLE_DEVICES", None) + env["ASCEND_RT_VISIBLE_DEVICES"] = visible_devices + env["VLLM_USE_AOT_COMPILE"] = "0" + cann_python_path = _cann_python_site_packages() + if cann_python_path is not None: + _prepend_pythonpath(env, cann_python_path) + if colocate: + env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" + return env + + +class NpuTrainingBootstrap(TrainingBootstrap): + def __init__(self) -> None: + self._bootstrapping = False + self._bootstrapped = False + + def bootstrap(self) -> None: + if self._bootstrapped or self._bootstrapping: + return + self._bootstrapping = True + try: + _ensure_torch_npu() + # Select NPU before MegatronAdaptor can make torch.cuda appear available. + register_npu_accelerator() + selected = accelerator.get_accelerator() + if selected.name != "npu": + raise RuntimeError(f"NPU bootstrap cannot use an already selected {selected.name!r} accelerator") + _install_safe_empty_cache() + # MegatronAdaptor must install its pre-patches before any Megatron module + # is imported. Apply the NPU attention override afterwards. + importlib.import_module("megatron_adaptor") + importlib.import_module("vime.backends.megatron_utils.npu_attention_patch") + except Exception: + # A failed bootstrap may be retried after the runtime environment is + # corrected; never leave a partially initialized success marker. + raise + else: + self._bootstrapped = True + finally: + self._bootstrapping = False + + def repatch(self, args: Any) -> None: + features_manager = importlib.import_module( + "megatron_adaptor.features_manager.features_manager" + ).FeaturesManager + full_args = importlib.import_module("megatron_adaptor.utils.args_utils").get_full_args() + for key, value in vars(args).items(): + setattr(full_args, key, value) + features_manager.remove_patches() + features_manager.apply_features_pre_patches(full_args) + features_manager.apply_features_patches(full_args) + # Repatch may replace attention again; importing a cached module alone + # does not reinstall Vime's existing override. + attention = importlib.import_module("vime.backends.megatron_utils.npu_attention_patch") + attention.DotProductAttention.forward = attention.npu_dot_product_attention_forward + + def adjust_tp_partition_dim(self, name: str, partition_dim: int) -> int: + if "linear_fc1.weight" in name or "linear_fc1.bias" in name: + return 0 + return partition_dim + + def training_context(self, offload_train: bool): + if not offload_train: + return nullcontext() + from torch_memory_saver import torch_memory_saver + + return torch_memory_saver.region(tag="training", enable_cpu_backup=True) + + def initialize_optimizer_state(self, optimizer: Any) -> None: + """Create lazy optimizer state before leaving the training memory pool.""" + if optimizer is None: + return + for opt in getattr(optimizer, "chained_optimizers", [optimizer]): + if opt.optimizer is not None and opt.init_state_fn is not None: + opt.init_state_fn(opt.optimizer, opt.config) + + +class NpuCheckpointCapabilities(CheckpointCapabilities): + def patch_default_planner(self, default_planner: Any) -> None: + if not hasattr(default_planner, "_validate_global_plan"): + return + + def _validate_global_plan(global_plan, metadata): + logger.info("[NPU checkpoint] Skipping validate_access_integrity") + return True + + default_planner._validate_global_plan = _validate_global_plan + + +def create_npu_platform() -> Platform: + # Ray workers also resolve the platform without importing Megatron. + register_npu_accelerator() + return Platform( + name="npu", + ray=NpuRayResourceSpec(), + weight_transfer=NpuWeightTransferPlatformOps(), + vllm=NpuVLLMLaunchPlatformOps(), + megatron=NpuTrainingBootstrap(), + checkpoint=NpuCheckpointCapabilities(default_megatron_to_hf_mode="bridge"), + ) diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py index 99f8922d5..2bb8ef5d6 100644 --- a/vime/ray/actor_group.py +++ b/vime/ray/actor_group.py @@ -1,11 +1,14 @@ import os +import shutil +import time +from pathlib import Path import ray from ray.util.placement_group import PlacementGroup from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST -from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath +from vime.platforms import current_platform +from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars class RayTrainGroup: @@ -35,14 +38,22 @@ def __init__( pg: tuple[PlacementGroup, list[int], list[int]], num_gpus_per_actor: float = 1, role: str = "actor", + with_ref: bool = False, + with_opd_teacher: bool = False, + actor_cls=None, ) -> None: self.args = args self._num_nodes = num_nodes self._num_gpus_per_node = num_gpus_per_node + self._pg = pg + self._num_gpus_per_actor = num_gpus_per_actor self.role = role - - # Allocate the GPUs for actors w/o instantiating them - self._allocate_gpus_for_actor(pg, num_gpus_per_actor) + self._actor_cls = actor_cls + self._with_ref = with_ref + self._with_opd_teacher = with_opd_teacher + self._rollout_manager = None + self._disk_weight_version = getattr(args, "update_weight_start_version", 0) + self._actor_handlers = [] def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): world_size = self._num_nodes * self._num_gpus_per_node @@ -60,20 +71,15 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): **self.args.train_env_vars, } + platform = current_platform() if self.args.offload_train and self.args.train_backend == "megatron": - import torch_memory_saver - - if is_npu(): - env_vars["TMS_HOOK_MODE"] = "torch" - env_vars["TMS_REGION_TAG"] = "training" - env_vars["TMS_ENABLE_CPU_BACKUP"] = "1" - if self.args.colocate: - env_vars["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" - cann_python_path = get_cann_python_site_packages() - if cann_python_path is not None: - prepend_pythonpath(env_vars, cann_python_path) + if platform.is_npu: + env_vars = platform.ray.train_runtime_env(self.args, env_vars) else: + import torch_memory_saver + for path in [ + "torch_memory_saver_hook_mode_preload_cu13.abi3.so", "torch_memory_saver_hook_mode_preload_cu12.abi3.so", "torch_memory_saver_hook_mode_preload.abi3.so", ]: @@ -83,52 +89,45 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): ) if os.path.exists(dynlib_path): break - else: - raise FileNotFoundError( - "Cannot find torch_memory_saver dynamic library. Please make sure torch_memory_saver is properly installed." - ) - - env_vars["LD_PRELOAD"] = dynlib_path - env_vars["TMS_INIT_ENABLE"] = "1" - env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" # We cannot do routing replay for critic. if self.args.use_routing_replay and self.role == "actor": env_vars["ENABLE_ROUTING_REPLAY"] = "1" - from vime.backends.megatron_utils.actor import MegatronTrainRayActor + if self._actor_cls is None: + from vime.backends.megatron_utils.actor import MegatronTrainRayActor - actor_impl = MegatronTrainRayActor + actor_impl = MegatronTrainRayActor + else: + actor_impl = self._actor_cls - TrainRayActor = ray.remote(runtime_env={"env_vars": env_vars})(actor_impl) - device_name = "NPU" if is_npu() else "GPU" + actor_options = { + "num_gpus": 1, + "runtime_env": {"env_vars": add_default_ray_env_vars(env_vars)}, + } + if getattr(self.args, "rollout_data_transport", "object-store") == "nixl": + actor_options["enable_tensor_transport"] = True + TrainRayActor = ray.remote(**actor_options)(actor_impl) # Create worker actors self._actor_handlers = [] master_addr, master_port = None, None for rank in range(world_size): + resource_options = {"num_gpus": num_gpus_per_actor} + if platform.is_npu: + resource_options = {"num_gpus": 0, **platform.ray.actor_options(num_gpus_per_actor)} actor = TrainRayActor.options( num_cpus=num_gpus_per_actor, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=reordered_bundle_indices[rank], ), - resources={device_name: num_gpus_per_actor}, + **resource_options, ).remote(world_size, rank, master_addr, master_port) if rank == 0: master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) self._actor_handlers.append(actor) - def async_init(self, args, role, with_ref=False, with_opd_teacher=False): - """ - Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. - """ - self.args = args - return [ - actor.init.remote(args, role, with_ref=with_ref, with_opd_teacher=with_opd_teacher) - for actor in self._actor_handlers - ] - def async_train(self, rollout_id, rollout_data_ref, external_data=None): """Do one rollout training. Returns a list of Ray refs (one per worker). @@ -151,11 +150,27 @@ def async_train(self, rollout_id, rollout_data_ref, external_data=None): def save_model(self, rollout_id, force_sync=False): """Save actor model""" - return ray.get([actor.save_model.remote(rollout_id, force_sync=force_sync) for actor in self._actor_handlers]) + ret = ray.get([actor.save_model.remote(rollout_id, force_sync=force_sync) for actor in self._actor_handlers]) + if self._release_train_enabled(): + self.args.load = self.args.save + self.args.ckpt_step = None + self.args.finetune = False + self.args.no_load_optim = self.args.no_save_optim + self.args.no_load_rng = False + return ret def update_weights(self): """Broadcast weights from rank 0 to all other ranks.""" - return ray.get([actor.update_weights.remote() for actor in self._actor_handlers]) + if not self._full_disk_weight_update_enabled(): + return ray.get([actor.update_weights.remote() for actor in self._actor_handlers]) + + weight_version = self._disk_weight_version + 1 + disk_weight_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{weight_version:06d}" + ray.get([actor.update_weights.remote() for actor in self._actor_handlers]) + self._disk_weight_version = weight_version + if self._release_train_enabled(): + self.release() + self._reload_rollout_weights_from_disk(disk_weight_dir, str(weight_version)) def onload(self): return ray.get([actor.wake_up.remote() for actor in self._actor_handlers]) @@ -163,8 +178,92 @@ def onload(self): def offload(self): return ray.get([actor.sleep.remote() for actor in self._actor_handlers]) + def release(self): + actors, self._actor_handlers = self._actor_handlers, [] + for actor in actors: + ray.kill(actor, no_restart=True) + if actors: + time.sleep(5) + + def create(self, rollout_manager=None): + if self._actor_handlers: + return None + if rollout_manager is not None: + self._rollout_manager = rollout_manager + self.args.update_weight_start_version = self._disk_weight_version + self._allocate_gpus_for_actor(self._pg, self._num_gpus_per_actor) + start_rollout_ids = ray.get( + [ + actor.init.remote( + self.args, + self.role, + with_ref=self._with_ref, + with_opd_teacher=self._with_opd_teacher, + ) + for actor in self._actor_handlers + ] + ) + if self._rollout_manager is not None: + self.set_rollout_manager(self._rollout_manager) + return start_rollout_ids + def clear_memory(self): return ray.get([actor.clear_memory.remote() for actor in self._actor_handlers]) def set_rollout_manager(self, rollout_manager): + self._rollout_manager = rollout_manager return ray.get([actor.set_rollout_manager.remote(rollout_manager) for actor in self._actor_handlers]) + + def _release_train_enabled(self): + return self.role == "actor" and getattr(self.args, "release_train", False) + + def _full_disk_weight_update_enabled(self): + return ( + self.role == "actor" + and self.args.update_weight_mode == "full" + and self.args.update_weight_transport == "disk" + ) + + def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version): + assert self._rollout_manager is not None, "disk weight update requires a rollout manager." + if self.args.offload_rollout: + ray.get(self._rollout_manager.onload_weights.remote()) + engines, *_ = ray.get(self._rollout_manager.get_updatable_engines_and_lock.remote()) + if not engines: + if not self.args.update_weight_disk_keep_files: + shutil.rmtree(disk_weight_dir, ignore_errors=True) + return + if self.args.update_weight_local_checkpoint_dir: + # each host pulls the published checkpoint onto local disk (e.g. NVMe) and + # the engines reload from there; the pull is disk-only, so it runs before + # pause and overlaps generation + ray.get([engine.pull_weights.remote(int(weight_version)) for engine in engines]) + model_path = self.args.update_weight_local_checkpoint_dir + else: + model_path = str(disk_weight_dir) + ray.get([engine.pause_generation.remote() for engine in engines]) + ray.get([engine.flush_cache.remote() for engine in engines]) + ray.get( + [ + engine.update_weights_from_disk.remote( + model_path=model_path, + weight_version=weight_version, + ) + for engine in engines + ] + ) + if self.args.ci_test: + engine_versions = ray.get([engine.get_weight_version.remote() for engine in engines]) + mismatches = [ + f"engine {idx}: {engine_version}" + for idx, engine_version in enumerate(engine_versions) + if str(engine_version) != str(weight_version) + ] + if mismatches: + raise RuntimeError( + "Weight version mismatch after disk reload! " + f"Expected: {weight_version}; " + ", ".join(mismatches) + ) + if not self.args.update_weight_disk_keep_files: + shutil.rmtree(disk_weight_dir, ignore_errors=True) + ray.get([engine.continue_generation.remote() for engine in engines]) diff --git a/vime/ray/placement_group.py b/vime/ray/placement_group.py index 97e735557..a08886e81 100644 --- a/vime/ray/placement_group.py +++ b/vime/ray/placement_group.py @@ -6,38 +6,29 @@ from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from vime.utils.common import is_npu +from vime.platforms import current_platform from .actor_group import RayTrainGroup -from .rollout import RolloutManager +from .utils import add_default_ray_env_vars logger = logging.getLogger(__name__) -# @ray.remote(num_gpus=1) -@ray.remote +@ray.remote(num_gpus=1) class InfoActor: def get_ip_and_gpu_id(self): - try: - import torch_npu # noqa: F401 - - has_npu = True - except ImportError: - has_npu = False + platform = current_platform() + if platform.is_npu: + accelerator_ids = platform.ray.accelerator_ids() + if accelerator_ids: + return ray.util.get_node_ip_address(), accelerator_ids[0] - if has_npu or is_npu(): - npu_ids = ray.get_runtime_context().get_accelerator_ids().get("NPU", []) - if npu_ids: - return ray.util.get_node_ip_address(), npu_ids[0] + raise RuntimeError( + f"No {platform.ray.resource_name} accelerator IDs found. " + f"Accelerator IDs: {ray.get_runtime_context().get_accelerator_ids()}" + ) - gpu_ids = ray.get_gpu_ids() - if gpu_ids: - return ray.util.get_node_ip_address(), gpu_ids[0] - - raise RuntimeError( - "No GPU/NPU IDs found. " - f"Accelerator IDs: {ray.get_runtime_context().get_accelerator_ids()}, GPU IDs: {gpu_ids}" - ) + return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0] def sort_key(x): @@ -63,22 +54,46 @@ def sort_key(x): def _create_placement_group(num_gpus): """Create a placement group with the specified number of GPUs.""" - device_name = "NPU" if is_npu() else "GPU" - bundles = [{device_name: 1, "CPU": 1} for _ in range(num_gpus)] + if num_gpus == 0: + return None, [], [] + + platform = current_platform() + resource_name = platform.ray.resource_name + bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] + if platform.is_npu: + bundles = [platform.ray.bundle_resources() for _ in range(num_gpus)] pg = placement_group(bundles, strategy="PACK") num_bundles = len(bundles) - ray.get(pg.ready()) + # Wait for the placement group to be scheduled. Poll rather than a bare + # ray.get(pg.ready()) so the wait is observable: when it can't be placed yet + # (a node's GPUs haven't registered with the GCS, or an autoscaler is still + # bringing nodes up) log the GPU counts periodically instead of hanging with no + # output. The wait stays unbounded, so autoscaling clusters — where a pending + # placement group is what drives scale-up — are unaffected. + ready_ref = pg.ready() + elapsed = 0 + log_interval = 30 + while not ray.wait([ready_ref], timeout=log_interval)[0]: + elapsed += log_interval + total = ray.cluster_resources().get(resource_name, 0) + available = ray.available_resources().get(resource_name, 0) + logger.info( + f"Waiting for placement group of {num_gpus} {resource_name} devices (elapsed {elapsed}s): " + f"{total:g} registered with Ray, {available:g} available." + ) + # use info actor to get the GPU id info_actors = [] for i in range(num_bundles): + resource_options = {"num_gpus": 0, **platform.ray.actor_options(1)} if platform.is_npu else {} info_actors.append( InfoActor.options( scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=i, ), - resources={device_name: 1}, + **resource_options, ).remote() ) gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors]) @@ -101,22 +116,30 @@ def _create_placement_group(num_gpus): return pg, pg_reordered_bundle_indices, pg_reordered_gpu_ids +def _get_placement_group_layout(args) -> tuple[int, int]: + actor_num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + + if args.debug_train_only: + return actor_num_gpus, 0 + + if args.rollout_external: + if args.debug_rollout_only: + return actor_num_gpus, 0 + return actor_num_gpus, actor_num_gpus + + if args.debug_rollout_only: + return args.rollout_num_gpus, 0 + + if args.colocate: + return max(actor_num_gpus, args.rollout_num_gpus), 0 + + return actor_num_gpus + args.rollout_num_gpus, actor_num_gpus + + def create_placement_groups(args): """Create placement groups for actor, critic, and rollout engines.""" - num_gpus = 0 - if args.debug_train_only: - num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node - rollout_offset = 0 - elif args.debug_rollout_only: - num_gpus = args.rollout_num_gpus - rollout_offset = 0 - elif args.colocate: - num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node - rollout_offset = 0 - else: - num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus - rollout_offset = args.actor_num_nodes * args.actor_num_gpus_per_node + num_gpus, rollout_offset = _get_placement_group_layout(args) logger.info(f"Creating placement group with {num_gpus} GPUs...") pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids = _create_placement_group(num_gpus) @@ -133,7 +156,16 @@ def create_placement_groups(args): return result -def allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="actor"): +def allocate_train_group( + args, + num_nodes, + num_gpus_per_node, + pg, + role="actor", + with_ref=False, + with_opd_teacher=False, + actor_cls=None, +): return RayTrainGroup( args=args, num_nodes=num_nodes, @@ -141,25 +173,40 @@ def allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="actor"): pg=pg, num_gpus_per_actor=0.4, role=role, + with_ref=with_ref, + with_opd_teacher=with_opd_teacher, + actor_cls=actor_cls, ) -def create_training_models(args, pgs, rollout_manager): +def create_actor_model(args, pgs, rollout_manager, actor_cls=None): actor_args = args if args.megatron_config_path is not None: from vime.utils.arguments import parse_megatron_role_args actor_args = parse_megatron_role_args(args, args.megatron_config_path, role="actor") + actor_model_kwargs = {} + if actor_cls is not None: + actor_model_kwargs["actor_cls"] = actor_cls actor_model = allocate_train_group( args=actor_args, num_nodes=args.actor_num_nodes, num_gpus_per_node=args.actor_num_gpus_per_node, pg=pgs["actor"], + with_ref=actor_args.kl_coef != 0 or actor_args.use_kl_loss, + with_opd_teacher=actor_args.use_opd and actor_args.opd_type == "megatron", + **actor_model_kwargs, ) + actor_start_rollout_ids = actor_model.create(rollout_manager=rollout_manager) + return actor_model, actor_start_rollout_ids + + +def create_training_models(args, pgs, rollout_manager, actor_cls=None): + actor_model, actor_start_rollout_ids = create_actor_model(args, pgs, rollout_manager, actor_cls=actor_cls) critic_model = None - if args.use_critic: + if args.use_critic and args.num_rollout != 0: from vime.utils.arguments import parse_megatron_role_args critic_args = ( @@ -177,18 +224,10 @@ def create_training_models(args, pgs, rollout_manager): pg=pgs["critic"], role="critic", ) - critic_start_rollout_ids = ray.get(critic_model.async_init(critic_model.args, role="critic", with_ref=False)) - - actor_start_rollout_ids = ray.get( - actor_model.async_init( - actor_args, - role="actor", - with_ref=actor_args.kl_coef != 0 or actor_args.use_kl_loss, - with_opd_teacher=actor_args.use_opd and actor_args.opd_type == "megatron", - ) - ) + critic_start_rollout_ids = critic_model.create(rollout_manager=rollout_manager) + # TODO how to decide rollout start id when critic is involved? For now we just require user to specify it via args. - if args.use_critic: + if critic_model is not None: start_rollout_ids = critic_start_rollout_ids else: start_rollout_ids = actor_start_rollout_ids @@ -198,10 +237,6 @@ def create_training_models(args, pgs, rollout_manager): if args.start_rollout_id is None: args.start_rollout_id = start_rollout_ids[0] - actor_model.set_rollout_manager(rollout_manager) - if args.use_critic: - critic_model.set_rollout_manager(rollout_manager) - if args.rollout_global_dataset: ray.get(rollout_manager.load.remote(args.start_rollout_id - 1)) @@ -209,12 +244,16 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): - device_name = "NPU" if is_npu() else "GPU" - rollout_manager = RolloutManager.options( - num_cpus=1, - # num_gpus=0, - resources={device_name: 0}, - ).remote(args, pg) + from .rollout import RolloutManager + + rollout_manager_options = { + "num_cpus": 1, + "num_gpus": 0, + "runtime_env": {"env_vars": add_default_ray_env_vars()}, + } + if getattr(args, "rollout_data_transport", "object-store") == "nixl": + rollout_manager_options["enable_tensor_transport"] = True + rollout_manager = RolloutManager.options(**rollout_manager_options).remote(args, pg) # calculate num_rollout from num_epoch num_rollout_per_epoch = None diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 0fe04e4a3..78d21d15f 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -1,38 +1,32 @@ -import dataclasses import itertools import logging -import multiprocessing -import random import time -from pathlib import Path from typing import Any -import numpy as np import ray import torch -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from vime.backends.vllm_utils.vllm_config import ModelConfig, ServerGroupConfig, VllmConfig -from vime.backends.vllm_utils.vllm_engine import VLLMEngine - -# Memory-type tag strings shared with the vLLM engine's sleep/wake_up API. -GPU_MEMORY_TYPE_KV_CACHE = "kv_cache" -GPU_MEMORY_TYPE_WEIGHTS = "weights" -GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" +from vime.backends.vllm_utils.deployment import start_rollout_servers +from vime.observability import logging_utils +from vime.observability.logging_utils import configure_logger, init_tracking +from vime.observability.rollout_data_utils import ( + load_debug_rollout_data, + save_debug_rollout_data, + tensorize_rollout_data_for_training, + validate_rollout_id_annotated, + validate_rollout_routed_experts_for_replay, +) +from vime.observability.rollout_metrics import log_eval_rollout_data, log_rollout_data from vime.rollout.base_types import call_rollout_fn -from vime.utils import logging_utils -from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath +from vime.rollout.sample_hooks import set_current_rollout_id +from vime.utils.data import get_source from vime.utils.dp_schedule import build_dp_schedule from vime.utils.health_monitor import RolloutHealthMonitor -from vime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client -from vime.utils.logging_utils import configure_logger, init_tracking -from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix -from vime.utils.misc import Box, group_by, load_function +from vime.utils.http_utils import init_http_client +from vime.utils.misc import Box, load_function from vime.utils.types import Sample -from ..utils.metric_utils import has_repetition -from .rollout_validation import validate_server_group_gpu_indices -from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock +from .utils import Lock, add_default_ray_env_vars logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) @@ -40,323 +34,6 @@ logger = logging.getLogger(__name__) -@dataclasses.dataclass -class ServerGroup: - """A group of homogeneous vLLM engines with the same configuration. - - All engines in a group share the same tp_size / nodes_per_engine / pg. - A RolloutServer may contain multiple ServerGroups (e.g. prefill vs decode - in PD disaggregation). - """ - - args: Any - pg: Any # (placement_group, reordered_bundle_indices, reordered_gpu_ids) - all_engines: list - num_gpus_per_engine: int - num_new_engines: int - worker_type: str = "regular" # "regular", "prefill", "decode", or "placeholder" - rank_offset: int = 0 # cumulative engine count before this group - gpu_offset: int = 0 # cumulative GPU count before this group - vllm_overrides: dict = dataclasses.field(default_factory=dict) - needs_offload: bool = False # True when this group's GPUs overlap with megatron - model_path: str | None = None # checkpoint path for update_weights_from_disk - router_ip: str | None = None - router_port: int | None = None - - @property - def nodes_per_engine(self): - return max(1, self.num_gpus_per_engine // self.args.num_gpus_per_node) - - @property - def engines(self): - """Node-0 engines only (for multi-node serving).""" - return self.all_engines[:: self.nodes_per_engine] - - def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[list, dict[int, int]]: - """Create Ray actors, allocate ports, and fire ``engine.init()`` without waiting. - - Returns ``(init_handles, port_cursors)`` where *init_handles* is a list - of Ray ObjectRefs and *port_cursors* maps node index → next free port. - The caller should ``ray.get()`` on the handles to block until the - engines are healthy, and pass *port_cursors* to the next server group - so that different groups on the same node don't race for ports. - - Placeholder groups (worker_type="placeholder") skip engine creation entirely. - """ - if port_cursors is None: - port_cursors = {} - if self.args.debug_train_only or self.worker_type == "placeholder": - self.num_new_engines = 0 - return [], port_cursors - - num_gpu_per_engine = min(self.num_gpus_per_engine, self.args.num_gpus_per_node) - - pg, reordered_bundle_indices, reordered_gpu_ids = self.pg - validate_server_group_gpu_indices( - worker_type=self.worker_type, - gpu_offset=self.gpu_offset, - num_gpus_per_engine=self.num_gpus_per_engine, - num_gpu_per_engine=num_gpu_per_engine, - num_engines=len(self.all_engines), - num_available_gpus=len(reordered_gpu_ids), - rollout_num_gpus=self.args.rollout_num_gpus, - rollout_num_gpus_per_engine=self.args.rollout_num_gpus_per_engine, - ) - - RolloutRayActor = ray.remote(VLLMEngine) - - rollout_engines = [] - device_name = "NPU" if is_npu() else "GPU" - for i in range(len(self.all_engines)): - if self.all_engines[i] is not None: - continue - - global_rank = self.rank_offset + i - num_gpus = 0.2 - num_cpus = num_gpus - - # Get the base GPU ID from placement group using gpu_offset. - gpu_index = self.gpu_offset + i * num_gpu_per_engine - base_gpu_id = int(reordered_gpu_ids[gpu_index]) - - scheduling_strategy = PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_capture_child_tasks=True, - placement_group_bundle_index=reordered_bundle_indices[gpu_index], - ) - - env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} - if is_npu(): - cann_python_path = get_cann_python_site_packages() - if cann_python_path is not None: - prepend_pythonpath(env_vars, cann_python_path) - if self.args.colocate: - env_vars["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False" - env_vars["VLLM_USE_AOT_COMPILE"] = "0" - rollout_engine = RolloutRayActor.options( - num_cpus=num_cpus, - scheduling_strategy=scheduling_strategy, - runtime_env={ - "env_vars": env_vars, - }, - resources={device_name: num_gpus}, - ).remote( - self.args, - rank=global_rank, - worker_type=self.worker_type, - base_gpu_id=base_gpu_id, - vllm_overrides=self.vllm_overrides, - num_gpus_per_engine=self.num_gpus_per_engine, - ) - - rollout_engines.append((global_rank, rollout_engine)) - self.all_engines[i] = rollout_engine - - self.num_new_engines = len(rollout_engines) - - if self.num_new_engines == 0: - return [], port_cursors - - if self.args.rollout_external: - addr_and_ports = _allocate_rollout_engine_addr_and_ports_external( - args=self.args, rollout_engines=rollout_engines - ) - else: - # Compute base_port from the maximum cursor across all nodes that - # this group's engines may land on (conservative: just use global max). - base_port = max(port_cursors.values()) if port_cursors else 15000 - addr_and_ports, port_cursors = _allocate_rollout_engine_addr_and_ports_normal( - args=self.args, - rollout_engines=rollout_engines, - worker_type=self.worker_type, - num_gpus_per_engine=self.num_gpus_per_engine, - rank_offset=self.rank_offset, - base_port=base_port, - ) - - init_handles = [ - engine.init.remote( - **(addr_and_ports[rank]), - router_ip=self.router_ip, - router_port=self.router_port, - ) - for rank, engine in rollout_engines - ] - return init_handles, port_cursors - - def offload(self): - """Fire release_memory_occupation on all engines (non-blocking). - - Returns a list of Ray ObjectRefs. Skipped for groups that do not - overlap with megatron GPUs (``needs_offload=False``). - """ - if not self.needs_offload: - return [] - return [engine.release_memory_occupation.remote() for engine in self.engines if engine is not None] - - def onload(self, tags: list[str] | None = None): - """Fire resume_memory_occupation on all engines (non-blocking). - - Returns a list of Ray ObjectRefs. Skipped for groups that do not - overlap with megatron GPUs (``needs_offload=False``). - """ - if not self.needs_offload: - return [] - return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.engines if engine is not None] - - def onload_weights_from_disk(self): - """Reload weights from ``model_path`` for non-updatable groups. - - Used instead of ``resume_memory_occupation(tags=[WEIGHTS])`` so that - CPU memory is not consumed by offloaded weight copies. - """ - if not self.needs_offload or not self.model_path: - return [] - return [ - engine.update_weights_from_disk.remote(self.model_path) for engine in self.engines if engine is not None - ] - - -@dataclasses.dataclass -class RolloutServer: - """A model served behind a shared router, with one or more server groups. - - Each RolloutServer represents one model deployed behind a single router. - A server may contain multiple ServerGroups with different - ``num_gpus_per_engine`` (e.g. prefill TP=2, decode TP=4). - """ - - server_groups: list[ServerGroup] - router_ip: str | None = None - router_port: int | None = None - prometheus_port: int | None = None - model_name: str = "default" - update_weights: bool = True - - @property - def engines(self): - """All node-0 engines across all groups (placeholder groups contribute nothing).""" - return [e for g in self.server_groups for e in g.engines] - - @property - def all_engines(self): - """All engines (including non-node-0) across all groups.""" - return [e for g in self.server_groups for e in g.all_engines] - - @property - def num_new_engines(self): - return sum(g.num_new_engines for g in self.server_groups) - - @num_new_engines.setter - def num_new_engines(self, value): - for g in self.server_groups: - g.num_new_engines = value - - @property - def engine_gpu_counts(self) -> list[int]: - """Per-engine GPU count for all node-0 engines, parallel to ``engines``.""" - return [g.num_gpus_per_engine for g in self.server_groups for _ in g.engines] - - @property - def engine_gpu_offsets(self) -> list[int]: - """Per-engine GPU offset for all node-0 engines, parallel to ``engines``. - - Accounts for placeholder groups that occupy GPU slots without creating engines. - """ - offsets = [] - for g in self.server_groups: - for j in range(len(g.engines)): - offsets.append(g.gpu_offset + j * g.num_gpus_per_engine) - return offsets - - @property - def nodes_per_engine(self): - """Nodes per engine. Only valid when all active groups share the same value.""" - values = {g.nodes_per_engine for g in self.server_groups if g.worker_type != "placeholder"} - if len(values) != 1: - raise ValueError(f"Heterogeneous nodes_per_engine across groups: {values}") - return values.pop() - - def recover(self): - """Recover dead engines across all active groups, overlapping init.""" - # Record dead indices per group before starting. - dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in self.server_groups] - - # Start all groups concurrently. - all_handles = [] - port_cursors: dict[int, int] = {} - for g in self.server_groups: - handles, port_cursors = g.start_engines(port_cursors) - all_handles.extend(handles) - if all_handles: - ray.get(all_handles) - - # Post-recovery: offload then onload weights for newly created engines. - release_handles = [] - updatable_new_engines = [] - non_updatable_groups_engines: list[tuple[str, list]] = [] - for g, dead_indices in zip(self.server_groups, dead_per_group, strict=True): - logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})") - assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length" - if g.needs_offload and dead_indices: - new_engines = [g.all_engines[i] for i in dead_indices] - release_handles.extend(engine.release_memory_occupation.remote() for engine in new_engines) - if self.update_weights: - updatable_new_engines.extend(new_engines) - elif g.model_path: - non_updatable_groups_engines.append((g.model_path, new_engines)) - - if release_handles: - ray.get(release_handles) - # Resume GPU memory for all engines that need offload. - all_resume_engines = updatable_new_engines[:] - for _model_path, engines in non_updatable_groups_engines: - all_resume_engines.extend(engines) - if all_resume_engines: - ray.get( - [ - engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - for engine in all_resume_engines - ] - ) - - def offload(self): - """Release memory occupation across all groups (concurrent).""" - handles = [] - for g in self.server_groups: - handles.extend(g.offload()) - return ray.get(handles) if handles else [] - - def onload(self, tags: list[str] | None = None): - """Resume memory occupation across all groups (concurrent).""" - handles = [] - for g in self.server_groups: - handles.extend(g.onload(tags)) - return ray.get(handles) if handles else [] - - def onload_weights(self): - """Restore weights for offloaded groups. - - All groups resume from CPU cache via ``resume_memory_occupation``. - For updatable servers, weights will be overwritten by - ``update_weights`` shortly after. For non-updatable servers the - CPU backup already contains the correct (unchanged) weights. - """ - handles = [] - for g in self.server_groups: - if not g.needs_offload: - continue - handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])) - return ray.get(handles) if handles else [] - - def onload_kv(self): - """Resume KV cache and CUDA graphs for offloaded groups.""" - handles = [] - for g in self.server_groups: - handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])) - return ray.get(handles) if handles else [] - - @ray.remote class RolloutManager: """The class to run rollout and convert rollout data to training data.""" @@ -367,6 +44,13 @@ def __init__(self, args, pg): self.pg = pg self.args = args + rollout_init_handles: list[Any] = [] + if self.args.debug_train_only: + self.servers: dict[str, Any] = {} + else: + init_http_client(args) + self.servers, rollout_init_handles = start_rollout_servers(args, pg) + data_source_cls = load_function(self.args.data_source_path) self.data_source = data_source_cls(args) @@ -383,15 +67,15 @@ def __init__(self, args, pg): logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.") - if self.args.debug_train_only: - self.servers: dict[str, RolloutServer] = {} - else: - init_http_client(args) - self.servers = start_rollout_servers(args, pg) + if rollout_init_handles: + ray.get(rollout_init_handles) init_tracking(args, primary=False) - device_name = "NPU" if is_npu() else "GPU" - self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0, resources={device_name: 0}).remote() + self.rollout_engine_lock = Lock.options( + num_cpus=1, + num_gpus=0, + runtime_env={"env_vars": add_default_ray_env_vars()}, + ).remote() self.rollout_id = -1 self._health_monitors = [] @@ -430,7 +114,12 @@ def _try_ci_fault_injection(self): # Only inject fault once self._ci_fault_injection_pending = False - if self.server and self.server.server_groups[0].all_engines and self.server.server_groups[0].all_engines[0]: + if ( + self.server + and self.server.server_groups + and self.server.server_groups[0].all_engines + and self.server.server_groups[0].all_engines[0] + ): logger.info("CI Fault Injection: Simulating crash on engine 0 during generate") try: # This will cause the ray actor to exit @@ -446,16 +135,19 @@ def _try_ci_fault_injection(self): def dispose(self): for monitor in self._health_monitors: monitor.stop() + engines = [engine for server in self.servers.values() for engine in server.all_engines if engine is not None] + if engines: + ray.get([engine.shutdown.remote() for engine in engines]) logging_utils.finish_tracking(self.args) @property - def server(self) -> RolloutServer | None: + def server(self) -> Any | None: """Default server (first model). For backward compatibility.""" if not self.servers: return None return next(iter(self.servers.values())) - def _get_updatable_server(self) -> RolloutServer | None: + def _get_updatable_server(self) -> Any | None: """Return the server with ``update_weights=True``. When multiple updatable servers exist, returns the first one @@ -482,8 +174,9 @@ def get_updatable_engines_and_lock(self): engines = srv.engines if srv else [] gpu_counts = srv.engine_gpu_counts if srv else [] gpu_offsets = srv.engine_gpu_offsets if srv else [] + parallel_configs = srv.engine_parallel_configs if srv else [] num_new = srv.num_new_engines if srv else 0 - return engines, self.rollout_engine_lock, num_new, gpu_counts, gpu_offsets + return engines, self.rollout_engine_lock, num_new, gpu_counts, gpu_offsets, parallel_configs def get_num_rollout_per_epoch(self): assert self.args.rollout_global_dataset @@ -492,12 +185,18 @@ def get_num_rollout_per_epoch(self): def generate(self, rollout_id): start_time = time.time() self.rollout_id = rollout_id + set_current_rollout_id(rollout_id) self.health_monitoring_resume() if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: self._try_ci_fault_injection() data, metrics = self._get_rollout_data(rollout_id=rollout_id) - self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False) - _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) + save_debug_rollout_data( + self.args.save_debug_rollout_data, + data, + rollout_id=rollout_id, + evaluation=False, + ) + log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) if self.args.debug_rollout_only: # if debug rollout only, we don't convert samples to train data and directly return return @@ -508,12 +207,18 @@ def eval(self, rollout_id): if self.args.debug_train_only: # if debug train only, we don't generate evaluation data return + set_current_rollout_id(rollout_id) self.health_monitoring_resume() result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True) data = result.data - self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True) - _log_eval_rollout_data(rollout_id, self.args, data, result.metrics) + save_debug_rollout_data( + self.args.save_debug_rollout_data, + data, + rollout_id=rollout_id, + evaluation=True, + ) + log_eval_rollout_data(rollout_id, self.args, data, result.metrics) def save(self, rollout_id): self.data_source.save(rollout_id) @@ -539,7 +244,7 @@ def onload_kv(self): srv.onload_kv() def recover_updatable_engines(self): - """Restart any dead rollout engines and update num_new_engines for update_weights detection. + """Restart dead updatable rollout engines before the next weight update. Recovers the updatable model (the one that receives weight updates from training). @@ -547,19 +252,9 @@ def recover_updatable_engines(self): self.health_monitoring_pause() srv = self._get_updatable_server() if self.rollout_id == -1 or srv is None: - engines = srv.engines if srv else [] - gpu_counts = srv.engine_gpu_counts if srv else [] - gpu_offsets = srv.engine_gpu_offsets if srv else [] - return engines, self.rollout_engine_lock, (srv.num_new_engines if srv else 0), gpu_counts, gpu_offsets + return srv.recover() - return ( - srv.engines, - self.rollout_engine_lock, - srv.num_new_engines, - srv.engine_gpu_counts, - srv.engine_gpu_offsets, - ) def clear_updatable_num_new_engines(self): # when fault tolerance is not enabled, we need to manually clear num_new_engines after update_weights @@ -580,18 +275,11 @@ def check_weights(self, action: str): def _get_rollout_data(self, rollout_id): if self.args.load_debug_rollout_data: - data = torch.load( - self.args.load_debug_rollout_data.format(rollout_id=rollout_id), - weights_only=False, - )["samples"] - data = [Sample.from_dict(sample) for sample in data] - if (ratio := self.args.load_debug_rollout_data_subsample) is not None: - original_num_rows = len(data) - rough_subsample_num_rows = int(original_num_rows * ratio) - data = data[: rough_subsample_num_rows // 2] + data[-rough_subsample_num_rows // 2 :] - logger.info( - f"Subsample loaded debug rollout data using {ratio=} and change num rows {original_num_rows} -> {len(data)}" - ) + data = load_debug_rollout_data( + self.args.load_debug_rollout_data, + rollout_id=rollout_id, + subsample_ratio=self.args.load_debug_rollout_data_subsample, + ) metrics = None else: data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False) @@ -603,39 +291,20 @@ def _get_rollout_data(self, rollout_id): # subagent paths that split one rollout into N training samples must # set the same rollout_id on every sibling so the loss reducer counts # the rollout once instead of N times. - _validate_rollout_id_annotated(data) + validate_rollout_id_annotated(data) # flatten the data if it is a list of lists while isinstance(data[0], list): data = list(itertools.chain.from_iterable(data)) return data, metrics - def _save_debug_rollout_data(self, data, rollout_id, evaluation: bool): - # TODO to be refactored (originally Buffer._set_data) - if (path_template := self.args.save_debug_rollout_data) is not None: - path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id))) - logger.info(f"Save debug rollout data to {path}") - path.parent.mkdir(parents=True, exist_ok=True) - - # TODO may improve the format - if evaluation: - dump_data = dict( - samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]] - ) - else: - dump_data = dict( - samples=[sample.to_dict() for sample in data], - ) - - torch.save(dict(rollout_id=rollout_id, **dump_data), path) - def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]): if self.custom_reward_post_process_func is not None: return self.custom_reward_post_process_func(self.args, samples) raw_rewards = [sample.get_reward_value(self.args) for sample in samples] if ( - self.args.advantage_estimator in ["grpo", "gspo", "reinforce_plus_plus_baseline"] + self.args.advantage_estimator in ["grpo", "gspo", "cispo", "reinforce_plus_plus_baseline"] and self.args.rewards_normalization ): # group norm @@ -648,7 +317,7 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]): mean = rewards.mean(dim=-1, keepdim=True) rewards = rewards - mean - if self.args.advantage_estimator in ["grpo", "gspo"] and self.args.grpo_std_normalization: + if self.args.advantage_estimator in ["grpo", "gspo", "cispo"] and self.args.grpo_std_normalization: std = rewards.std(dim=-1, keepdim=True) rewards = rewards / (std + 1e-6) @@ -668,15 +337,15 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl assert len(raw_rewards) == len(samples) assert len(rewards) == len(samples) - # Rollout id (one per rollout execution). Default rollouts emit one - # sample per rollout, so we fall back to ``sample.index`` (unique). - # Compact / subagent paths that emit multiple training samples per - # rollout set ``rollout_id`` explicitly so all siblings share a - # value; the loss reducer then aggregates them as one rollout. - if samples[0].rollout_id is None: - rollout_ids = list(range(len(samples))) - else: - rollout_ids = [sample.rollout_id for sample in samples] + rollout_ids = [sample.rollout_id for sample in samples] + existed_rollout_id_values = set(rid for rid in rollout_ids if rid is not None) + tmp_id = 0 + for i in range(len(rollout_ids)): + if rollout_ids[i] is None: + while tmp_id in existed_rollout_id_values: + tmp_id += 1 + rollout_ids[i] = tmp_id + existed_rollout_id_values.add(tmp_id) train_data = { "tokens": [sample.tokens for sample in samples], @@ -739,8 +408,27 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].rollout_log_probs is not None: train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] + if getattr(self.args, "rollout_top_p", 1.0) != 1.0 and samples[0].rollout_top_p_token_ids is not None: + for sample in samples: + assert sample.rollout_top_p_token_ids is not None + assert sample.rollout_top_p_token_offsets is not None + assert len(sample.rollout_top_p_token_offsets) == sample.response_length + 1, ( + f"top-p token offsets length {len(sample.rollout_top_p_token_offsets)} " + f"!= response length + 1 {sample.response_length + 1}" + ) + offset_end = int(sample.rollout_top_p_token_offsets[-1]) + assert offset_end == len(sample.rollout_top_p_token_ids), ( + f"top-p token offsets[-1] {offset_end} " + f"!= token ids length {len(sample.rollout_top_p_token_ids)}" + ) + train_data["rollout_top_p_token_ids"] = [sample.rollout_top_p_token_ids for sample in samples] + train_data["rollout_top_p_token_offsets"] = [sample.rollout_top_p_token_offsets for sample in samples] + if samples[0].rollout_routed_experts is not None: - train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] + routed_experts = [torch.as_tensor(sample.rollout_routed_experts) for sample in samples] + if getattr(self.args, "use_rollout_routing_replay", False): + validate_rollout_routed_experts_for_replay(routed_experts, self.args) + train_data["rollout_routed_experts"] = routed_experts if samples[0].train_metadata is not None: train_data["metadata"] = [sample.train_metadata for sample in samples] @@ -751,6 +439,9 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].teacher_log_probs is not None: train_data["teacher_log_probs"] = [sample.teacher_log_probs for sample in samples] + if samples[0].metadata is not None: + train_data["source_names"] = [get_source(sample) for sample in samples] + return train_data def set_train_parallel_config(self, config: dict): @@ -797,7 +488,10 @@ def _split_train_data_by_dp(self, data): "rollout_ids", "rollout_mask_sums", "rollout_log_probs", + "rollout_top_p_token_ids", + "rollout_top_p_token_offsets", "rollout_routed_experts", + "source_names", "prompt", "teacher_log_probs", ]: @@ -812,551 +506,12 @@ def _split_train_data_by_dp(self, data): rollout_data["global_batch_sizes"] = global_batch_sizes rollout_data["num_microbatches"] = num_microbatches rollout_data["micro_batch_indices"] = micro_batch_indices[r] - rollout_data_refs.append(Box(ray.put(rollout_data))) + tensorize_rollout_data_for_training(rollout_data) + transport = getattr(self.args, "rollout_data_transport", "object-store") + if transport == "nixl": + rollout_data_refs.append(Box(ray.put(rollout_data, _tensor_transport="nixl"))) + elif transport == "object-store": + rollout_data_refs.append(Box(ray.put(rollout_data))) + else: + raise ValueError(f"Unsupported rollout data transport: {transport!r}") return rollout_data_refs - - -def _validate_rollout_id_annotated(node, depth=0): - """Walk the rollout function's nested output and validate ``rollout_id`` only - when a compact / subagent pattern is detected. - - "Compact" = the rollout function wraps multiple training samples from one - rollout execution into a ``list[Sample]``. In vime's convention the - default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout) - so its leaf ``list[Sample]`` lands at depth 1 and we skip validation, - preserving backward compatibility. A compact rollout adds a third level: - ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout), - so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require - every sibling to carry a non-None ``rollout_id`` and to share the same - value, so the loss reducer counts the rollout once instead of N times. - """ - if isinstance(node, Sample): - return - assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}" - if node and isinstance(node[0], Sample): - if depth >= 2 and len(node) > 1: - rids = [s.rollout_id for s in node] - missing = [i for i, r in enumerate(rids) if r is None] - assert not missing, ( - f"Compact rollout returned {len(node)} samples but rollout_id is unset on " - f"positions {missing}. Set Sample.rollout_id on every sibling so the loss " - "reducer can aggregate them as one rollout instead of N." - ) - assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}." - return - for item in node: - _validate_rollout_id_annotated(item, depth + 1) - - -def _allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): - addr_and_ports = {} - for rank, _ in rollout_engines: - addr = args.rollout_external_engine_addrs[rank] - [host, port] = addr.split(":") - addr_and_ports[rank] = dict( - dist_init_addr=addr, - nccl_port=None, - host=host, - port=int(port), - ) - return addr_and_ports - - -def _allocate_rollout_engine_addr_and_ports_normal( - *, - args, - rollout_engines, - worker_type="regular", - num_gpus_per_engine=None, - rank_offset=0, - base_port=15000, -): - # get ports - # there are 4 ports we need to allocate - # 1. server port - # 2. nccl port - # 3. dist_init_addr port - # 4. other ports for dp_attention, which is of size 4 + dp_size - _gpus_per_engine = num_gpus_per_engine or args.rollout_num_gpus_per_engine - num_engines_per_node = max(1, args.num_gpus_per_node // _gpus_per_engine) - addr_and_ports: dict[int, dict] = {} - - # Track per-node port cursors so that different server groups (called - # sequentially) never race for the same ports on a given node. - node_port_cursor: dict[int, int] = {} - - visited_nodes = set() - for rank, engine in rollout_engines: - local_rank = rank - rank_offset - node_index = local_rank // num_engines_per_node - if node_index in visited_nodes: - continue - visited_nodes.add(node_index) - # TODO: currently when restarting engines, we will set port for all engines on this node starting with this rank. - # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node. - num_engines_on_this_node = num_engines_per_node - (local_rank % num_engines_per_node) - - def get_addr_and_ports(engine, node_idx): - # use small ports to prevent ephemeral port between 32768 and 65536. - # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition - start_port = node_port_cursor.get(node_idx, base_port) - - def port(consecutive=1): - nonlocal start_port - _, port = ray.get( - engine._get_current_node_ip_and_free_port.remote( - start_port=start_port, - consecutive=consecutive, - ) - ) - start_port = port + consecutive - node_port_cursor[node_idx] = start_port - return port - - def addr(): - addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) - return addr - - return addr, port - - get_addr, get_port = get_addr_and_ports(engine, node_index) - - for i in range(num_engines_on_this_node): - current_rank = rank + i - addr_and_ports.setdefault(current_rank, {}) - addr_and_ports[current_rank]["host"] = get_addr() - addr_and_ports[current_rank]["port"] = get_port() - addr_and_ports[current_rank]["nccl_port"] = get_port() - - if worker_type in ("prefill", "decode"): - addr_and_ports[current_rank]["disaggregation_bootstrap_port"] = get_port() - - if _gpus_per_engine > args.num_gpus_per_node: - num_node_per_engine = _gpus_per_engine // args.num_gpus_per_node - if local_rank % num_node_per_engine == 0: - # this is the first node in the engine, we need to allocate the dist_init_addr port - dist_init_addr = f"{get_addr()}:{get_port(30 + args.vllm_dp_size)}" - for i in range(num_node_per_engine): - addr_and_ports.setdefault(rank + i, {}) - addr_and_ports[rank + i]["dist_init_addr"] = dist_init_addr - else: - for i in range(num_engines_on_this_node): - addr_and_ports[rank + i]["dist_init_addr"] = f"{get_addr()}:{get_port(30 + args.vllm_dp_size)}" - - for i, _ in rollout_engines: - for key in ["port", "nccl_port", "dist_init_addr"]: - assert key in addr_and_ports[i], f"Engine {i} {key} is not set." - logger.info(f"Ports for engine {i}: {addr_and_ports[i]}") - - return addr_and_ports, node_port_cursor - - -def _start_router( - args, - *, - has_pd_disaggregation: bool = False, - force_new: bool = False, - bind: tuple[str, int] | None = None, - prefill_urls: list | None = None, - decode_urls: list | None = None, -) -> tuple[str, int, int]: - """Start the rollout HTTP gateway (vllm-router).""" - if bind is not None: - router_ip, router_port = bind - else: - if not force_new and args.vllm_router_ip is not None: - return args.vllm_router_ip, args.vllm_router_port, None - router_ip = _wrap_ipv6(get_host_info()[1]) - if force_new or args.vllm_router_port is None: - router_port = find_available_port(random.randint(3000, 4000)) - else: - router_port = args.vllm_router_port - - from vllm_router.router_args import RouterArgs - - from vime.utils.http_utils import run_router - - router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) - router_args.host = router_ip - router_args.port = router_port - router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) - router_args.log_level = "warning" - router_args.request_timeout_secs = args.router_request_timeout_secs - - if has_pd_disaggregation: - router_args.vllm_pd_disaggregation = True - # Disable circuit breaker so transient RDMA transfer timeouts (PCIe - # contention under load) don't mark decode workers dead. - router_args.disable_circuit_breaker = True - - if prefill_urls is not None: - router_args.prefill_urls = prefill_urls - router_args.decode_urls = decode_urls - - logger.info(f"Launch router with args: {router_args}") - - process = multiprocessing.Process(target=run_router, args=(router_args,)) - process.daemon = True - process.start() - time.sleep(3) - assert process.is_alive() - logger.info(f"Router launched at {router_ip}:{router_port}, Prometheus port: {router_args.prometheus_port}") - return router_ip, router_port, router_args.prometheus_port - - -def _compute_rollout_offset(args) -> int: - """Offset (in PG bundle slots) where rollout GPUs start.""" - if args.debug_train_only or args.debug_rollout_only or args.colocate: - return 0 - offset = args.actor_num_nodes * args.actor_num_gpus_per_node - return offset - - -def _compute_megatron_num_gpus(args) -> int: - """Total number of megatron (actor + critic) GPU slots in the placement group.""" - if args.debug_rollout_only: - return 0 - num = args.actor_num_nodes * args.actor_num_gpus_per_node - return num - - -def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: - """Start rollout servers: one per model, each with its own router. - - Each model defined in the vLLM config gets its own router and set - of server groups. Server groups within a model may have different - ``num_gpus_per_engine`` (e.g. for PD disaggregation where prefill - and decode use different TP sizes). - - Returns a dict mapping model name → ``RolloutServer``. - - Note: ``init_http_client`` should be called separately before this, - as the HTTP client is shared across all servers. - """ - config = _resolve_vllm_config(args) - - servers: dict[str, RolloutServer] = {} - gpu_offset = 0 - engine_offset = 0 - - # Compute megatron GPU range for per-group offload decisions. - rollout_pg_offset = _compute_rollout_offset(args) - megatron_num_gpus = _compute_megatron_num_gpus(args) - - for model_idx, model_cfg in enumerate(config.models): - model_cfg.resolve(args) - - has_pd = model_cfg.has_pd_disaggregation - use_static_pd_router = has_pd - if use_static_pd_router: - router_ip = _wrap_ipv6(get_host_info()[1]) - router_port = find_available_port(random.randint(3000, 4000)) - prom_port = None # assigned when the router actually launches, after URL collection - engine_router_ip, engine_router_port = None, None - else: - router_ip, router_port, prom_port = _start_router( - args, has_pd_disaggregation=has_pd, force_new=(model_idx > 0) - ) - engine_router_ip, engine_router_port = router_ip, router_port - - # Write back so downstream readers (vllm_rollout, vllm_engine) see the - # router we just started (only relevant for first model in multi-model setups). - if model_idx == 0: - args.vllm_router_ip = router_ip - args.vllm_router_port = router_port - - server_groups: list[ServerGroup] = [] - port_cursors: dict[int, int] = {} - - has_epd = model_cfg.has_encoder_disaggregation - - def _make_group(group_cfg, router_ip, router_port, overrides_extra=None): - nonlocal engine_offset, gpu_offset - gpus_per_engine = group_cfg.num_gpus_per_engine - num_gpu_per_engine_local = min(gpus_per_engine, args.num_gpus_per_node) - num_engines = group_cfg.num_gpus // num_gpu_per_engine_local - - group_abs_start = rollout_pg_offset + gpu_offset - needs_offload = args.offload_rollout and group_abs_start < megatron_num_gpus - overrides = dict(group_cfg.overrides) - if overrides_extra: - for k, v in overrides_extra.items(): - overrides.setdefault(k, v) - if args.offload_rollout and not needs_offload: - overrides.setdefault("enable_memory_saver", False) - logger.info( - f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} " - f"(abs={group_abs_start}): needs_offload={needs_offload}" - ) - - group = ServerGroup( - args=args, - pg=pg, - all_engines=[None] * num_engines if group_cfg.worker_type != "placeholder" else [], - num_gpus_per_engine=gpus_per_engine, - num_new_engines=0, - worker_type=group_cfg.worker_type, - rank_offset=engine_offset, - gpu_offset=gpu_offset, - vllm_overrides=overrides, - needs_offload=needs_offload, - model_path=overrides.get("model_path", args.hf_checkpoint), - router_ip=router_ip, - router_port=router_port, - ) - engine_offset += num_engines - gpu_offset += group_cfg.num_gpus - return group - - if has_epd: - # --- Phase 1: start encoder groups, wait, collect URLs --- - encoder_urls: list[str] = [] - for group_cfg in model_cfg.server_groups: - if group_cfg.worker_type != "encoder": - continue - group = _make_group(group_cfg, engine_router_ip, engine_router_port) - handles, port_cursors = group.start_engines(port_cursors) - if handles: - ray.get(handles) - urls = ray.get([e.get_url.remote() for e in group.engines]) - encoder_urls.extend(u for u in urls if u is not None) - server_groups.append(group) - - logger.info(f"EPD phase 1 done: collected {len(encoder_urls)} encoder URLs: {encoder_urls}") - - # --- Phase 2: start non-encoder groups, injecting encoder URLs into - # language-only LLM workers. Prefill groups use this for full EPD, - # while regular groups allow encoder/LLM split without PD. - non_encoder_handles: list = [] - for group_cfg in model_cfg.server_groups: - if group_cfg.worker_type == "encoder": - continue - overrides_extra = {} - if encoder_urls and group_cfg.worker_type in ("prefill", "regular"): - overrides_extra["language_only"] = True - overrides_extra["encoder_urls"] = encoder_urls - group = _make_group(group_cfg, engine_router_ip, engine_router_port, overrides_extra=overrides_extra) - handles, port_cursors = group.start_engines(port_cursors) - non_encoder_handles.extend(handles) - server_groups.append(group) - - if non_encoder_handles: - ray.get(non_encoder_handles) - else: - # No EPD — start all groups in one pass (original path). - all_init_handles: list = [] - for group_cfg in model_cfg.server_groups: - group = _make_group(group_cfg, engine_router_ip, engine_router_port) - handles, port_cursors = group.start_engines(port_cursors) - all_init_handles.extend(handles) - server_groups.append(group) - - if all_init_handles: - ray.get(all_init_handles) - - if use_static_pd_router: - prefill_urls: list[tuple] = [] - decode_urls: list[str] = [] - for g in server_groups: - for e in g.engines: - if e is None: - continue - if g.worker_type == "prefill": - url = ray.get(e.get_url.remote()) - if url: - prefill_urls.append((url, None)) - elif g.worker_type == "decode": - url = ray.get(e.get_url.remote()) - if url: - decode_urls.append(url) - _, _, prom_port = _start_router( - args, - has_pd_disaggregation=True, - bind=(router_ip, router_port), - prefill_urls=prefill_urls, - decode_urls=decode_urls, - ) - - servers[model_cfg.name] = RolloutServer( - server_groups=server_groups, - router_ip=router_ip, - router_port=router_port, - model_name=model_cfg.name, - update_weights=model_cfg.update_weights, - prometheus_port=prom_port, - ) - - # Expose per-model router info for custom rollout functions. - args.vllm_model_routers = {name: (srv.router_ip, srv.router_port) for name, srv in servers.items()} - - return servers - - -def _resolve_vllm_config(args) -> VllmConfig: - """Build a VllmConfig from args, choosing the right source.""" - if getattr(args, "vllm_config", None): - config = VllmConfig.from_yaml(args.vllm_config) - # Validate total GPUs match. - expected = args.rollout_num_gpus - actual = config.total_num_gpus - assert actual == expected, f"vllm_config total GPUs ({actual}) != rollout_num_gpus ({expected})" - return config - - if args.prefill_num_servers is not None: - return VllmConfig.from_prefill_num_servers(args) - - # Default: single regular group. - return VllmConfig( - models=[ - ModelConfig( - name="default", - server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], - ) - ] - ) - - -def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): - if args.custom_eval_rollout_log_function_path is not None: - custom_log_func = load_function(args.custom_eval_rollout_log_function_path) - if custom_log_func(rollout_id, args, data, extra_metrics): - return - - log_dict = extra_metrics or {} - for key in data.keys(): - rewards = data[key]["rewards"] - log_dict[f"eval/{key}"] = sum(rewards) / len(rewards) - if (samples := data[key].get("samples")) is not None: - log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), f"eval/{key}/") - if "truncated" in data[key]: - truncated = data[key]["truncated"] - log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated) - if args.log_passrate: - log_dict |= dict_add_prefix( - compute_pass_rate( - flat_rewards=rewards, - group_size=args.n_samples_per_eval_prompt, - ), - f"eval/{key}-", - ) - - logger.info(f"eval {rollout_id}: {log_dict}") - - step = compute_rollout_step(args, rollout_id) - log_dict["eval/step"] = step - logging_utils.log(args, log_dict, step_key="eval/step") - - return log_dict - - -def _log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): - if args.custom_rollout_log_function_path is not None: - custom_log_func = load_function(args.custom_rollout_log_function_path) - if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time): - return - - if args.load_debug_rollout_data: - return - - log_dict = {**(rollout_extra_metrics or {})} - log_dict |= dict_add_prefix(compute_metrics_from_samples(args, samples), "rollout/") - log_dict |= dict_add_prefix(compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/") - logger.info(f"perf {rollout_id}: {log_dict}") - step = compute_rollout_step(args, rollout_id) - log_dict["rollout/step"] = step - logging_utils.log(args, log_dict, step_key="rollout/step") - - -def compute_metrics_from_samples(args, samples): - response_lengths = [sample.effective_response_length for sample in samples] - - log_dict = {} - log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/") - log_dict |= _compute_zero_std_metrics(args, samples) - log_dict |= _compute_spec_metrics(args, samples) - log_dict |= _compute_prefix_cache_metrics(args, samples) - log_dict |= _compute_reward_cat_metrics(args, samples) - log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item() - log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item() - return log_dict - - -def compute_perf_metrics_from_samples(args, samples, rollout_time): - non_generation_time = [sample.non_generation_time for sample in samples] - - log_dict = {} - log_dict["rollout_time"] = rollout_time - if max(non_generation_time) > 0: - log_dict |= dict_add_prefix(compute_statistics(non_generation_time), "non_generation_time/") - - def token_perf(response_lengths, non_generation_time, key=""): - max_response_length = max(response_lengths) - if args.rollout_num_gpus: - log_dict[f"{key}tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus - log_dict[f"longest_{key}sample_tokens_per_sec"] = max_response_length / rollout_time - - if max(non_generation_time) == 0: - return - - non_generation_time = [ - t for t, length in zip(non_generation_time, response_lengths, strict=True) if length == max_response_length - ] - mean_non_generation_time = sum(non_generation_time) / len(non_generation_time) - - log_dict[f"longest_{key}sample_non_generation_time"] = mean_non_generation_time - log_dict[f"longest_{key}sample_tokens_per_sec_without_non_generation"] = max_response_length / ( - rollout_time - mean_non_generation_time - ) - - token_perf([sample.response_length for sample in samples], non_generation_time, key="") - token_perf([sample.effective_response_length for sample in samples], non_generation_time, key="effective_") - - return log_dict - - -def _compute_zero_std_metrics(args, all_samples: list[Sample]): - # only compute in GRPO-like algorithms where one prompt has multiple responses - if args.advantage_estimator == "ppo": - return {} - - def _is_zero_std(samples: list[Sample]): - rewards = [sample.get_reward_value(args) for sample in samples] - return len(rewards) == 0 or all(rewards[0] == r for r in rewards) - - all_sample_groups = group_by(all_samples, lambda s: s.group_index) - interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)] - - interesting_rewards = [str(round(g[0].get_reward_value(args), 1)) for g in interesting_sample_groups] - - return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} - - -def _compute_spec_metrics(args, all_samples: list[Sample]): - if getattr(args, "vllm_speculative_config", None) is None: - return {} - num_samples = len(all_samples) - metrics = {} - metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples - metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples - return metrics - - -def _compute_prefix_cache_metrics(args, all_samples: list[Sample]): - num_samples = len(all_samples) - metrics = {} - total_cached_tokens = sum(sample.prefix_cache_info.cached_tokens for sample in all_samples) - total_prompt_tokens = sum(sample.prefix_cache_info.total_prompt_tokens for sample in all_samples) - - metrics["prefix_cache_hit_rate"] = total_cached_tokens / total_prompt_tokens if total_prompt_tokens > 0 else 0.0 - metrics["avg_cached_tokens_per_sample"] = total_cached_tokens / num_samples - return metrics - - -def _compute_reward_cat_metrics(args, all_samples: list[Sample]): - reward_cat_key = args.log_reward_category - if reward_cat_key is None: - return {} - - samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key]) - - return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()} diff --git a/vime/ray/rollout_validation.py b/vime/ray/rollout_validation.py deleted file mode 100644 index f27a7c172..000000000 --- a/vime/ray/rollout_validation.py +++ /dev/null @@ -1,32 +0,0 @@ -def validate_server_group_gpu_indices( - *, - worker_type: str, - gpu_offset: int, - num_gpus_per_engine: int, - num_gpu_per_engine: int, - num_engines: int, - num_available_gpus: int, - rollout_num_gpus: int, - rollout_num_gpus_per_engine: int, -) -> None: - if num_engines == 0: - return - - required_gpu_slots = gpu_offset + num_engines * num_gpu_per_engine - if gpu_offset >= 0 and num_gpu_per_engine > 0 and required_gpu_slots <= num_available_gpus: - return - - raise ValueError( - "Invalid rollout server group GPU placement: " - f"worker_type={worker_type}, " - f"gpu_offset={gpu_offset}, " - f"num_gpus_per_engine={num_gpus_per_engine}, " - f"num_gpu_per_engine_on_node={num_gpu_per_engine}, " - f"num_engines={num_engines}, " - f"required_gpu_slots={required_gpu_slots}, " - f"len(reordered_gpu_ids)={num_available_gpus}, " - f"rollout_num_gpus={rollout_num_gpus}, " - f"rollout_num_gpus_per_engine={rollout_num_gpus_per_engine}. " - "Please align --rollout-num-gpus, --rollout-num-gpus-per-engine, " - "and --vllm-config server_groups." - ) diff --git a/vime/ray/train_actor.py b/vime/ray/train_actor.py index 6da82343b..a9da1b065 100644 --- a/vime/ray/train_actor.py +++ b/vime/ray/train_actor.py @@ -9,27 +9,22 @@ import torch.distributed as dist import vime.utils.eval_config +from vime.observability.logging_utils import configure_logger +from vime.platforms import current_platform from vime.ray.ray_actor import RayActor -from vime.utils.common import is_npu +from vime.utils import accelerator from vime.utils.distributed_utils import init_gloo_group -from vime.utils.logging_utils import configure_logger from vime.utils.memory_utils import clear_memory, print_memory logger = logging.getLogger(__name__) def get_local_gpu_id(): - if is_npu(): - env_var = "ASCEND_RT_VISIBLE_DEVICES" - device_ids = ray.get_runtime_context().get_accelerator_ids()["NPU"] - else: - env_var = "CUDA_VISIBLE_DEVICES" - device_ids = ray.get_gpu_ids() - cvd = os.environ.get(env_var, None) - if cvd is None: - return device_ids[0] - else: - return cvd.split(",").index(str(device_ids[0])) + platform = current_platform() + if platform.is_npu: + return platform.ray.local_device_id() + + return accelerator.resolve_visible_device_id(ray.get_gpu_ids()[0]) class TrainRayActor(RayActor): @@ -63,12 +58,14 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): torch.serialization.add_safe_globals([vime.utils.eval_config.EvalDatasetConfig]) local_rank = int(os.environ.get("LOCAL_RANK", 0)) - if is_npu(): - torch.npu.set_device(f"npu:{local_rank}") - else: - torch.cuda.set_device(f"cuda:{local_rank}") + accelerator.set_device(local_rank) + if accelerator.set_allocator_expandable_segments(): + logger.info( + f"[Rank {self._rank}] Enabled {accelerator.device_type().upper()} memory allocator " + "expandable_segments for train actor" + ) - backend = args.distributed_backend + backend = accelerator.process_group_backend(args.distributed_backend) dist.init_process_group( backend=backend, @@ -80,17 +77,20 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): args.world_size = dist.get_world_size() try: - import pynvml + if torch.version.hip is not None: + logger.info("Detected ROCm/HIP environment, skipping NUMA affinity setup") + else: + import pynvml - pynvml.nvmlInit() + pynvml.nvmlInit() - local_rank = int(os.environ["RANK"]) % args.num_gpus_per_node + local_rank = int(os.environ["RANK"]) % args.num_gpus_per_node - handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank) - pynvml.nvmlDeviceSetCpuAffinity(handle) + handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank) + pynvml.nvmlDeviceSetCpuAffinity(handle) - logger.info(f"Set NUMA affinity for GPU {local_rank}") - pynvml.nvmlShutdown() + logger.info(f"Set NUMA affinity for GPU {local_rank}") + pynvml.nvmlShutdown() except ImportError: logger.info("Warning: pynvml not available, skipping NUMA affinity setup") @@ -124,10 +124,6 @@ def save_model(self, rollout_id, force_sync=False): def update_weights(self): raise NotImplementedError - @abc.abstractmethod - def _get_parallel_config(self): - raise NotImplementedError - def set_rollout_manager(self, rollout_manager): self.rollout_manager = rollout_manager if not self.args.debug_rollout_only and self.args.rank == 0: diff --git a/vime/ray/utils.py b/vime/ray/utils.py index b4103f5c4..c3bc64537 100644 --- a/vime/ray/utils.py +++ b/vime/ray/utils.py @@ -2,9 +2,9 @@ import os import ray -import torch -from vime.ray.ray_actor import RayActor +from vime.ray.ray_actor import RayActor +from vime.utils import accelerator # Refer to # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96 @@ -16,6 +16,7 @@ # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98 NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", @@ -24,15 +25,24 @@ "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", ] +RAY_DEFAULT_ENV_VARS = { + # Ray's uvloop integration has caused intermittent async actor issues. + "RAY_USE_UVLOOP": "0", +} + + +def add_default_ray_env_vars(env_vars: dict[str, str] | None = None) -> dict[str, str]: + return RAY_DEFAULT_ENV_VARS | (env_vars or {}) + def ray_noset_visible_devices(env_vars=os.environ): return any(env_vars.get(env_var) for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST) def get_physical_gpu_id(): - device = torch.cuda.current_device() - props = torch.cuda.get_device_properties(device) - return str(props.uuid) + device = accelerator.current_device() + props = accelerator.get_device_properties(device) + return str(getattr(props, "uuid", device)) @ray.remote diff --git a/vime/rollout/filter_hub/base_types.py b/vime/rollout/filter_hub/base_types.py index 2937273bd..5f2154c2c 100644 --- a/vime/rollout/filter_hub/base_types.py +++ b/vime/rollout/filter_hub/base_types.py @@ -6,6 +6,22 @@ class DynamicFilterOutput: keep: bool reason: str | None = None + # Keep a rejected group when dropping it would leave too few candidates to + # fill the rollout batch. This avoids launching another oversampling round. + keep_when_insufficient: bool = False + + +def should_drop_dynamic_filter_output( + output: DynamicFilterOutput, + *, + remaining_batch_size: int, + target_data_size: int, +) -> bool: + if output.keep: + return False + if output.keep_when_insufficient and remaining_batch_size <= target_data_size: + return False + return True def call_dynamic_filter(fn, *args, **kwargs): diff --git a/vime/rollout/filter_hub/dynamic_sampling_filters.py b/vime/rollout/filter_hub/dynamic_sampling_filters.py index 743c420e2..df88f866a 100644 --- a/vime/rollout/filter_hub/dynamic_sampling_filters.py +++ b/vime/rollout/filter_hub/dynamic_sampling_filters.py @@ -3,7 +3,7 @@ from vime.rollout.filter_hub.base_types import DynamicFilterOutput from vime.utils.types import Sample -__all__ = ["check_reward_nonzero_std"] +__all__ = ["check_reward_nonzero_std", "check_reward_nonzero_std_with_fallback"] def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): @@ -13,3 +13,11 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): keep=keep, reason=None if keep else f"zero_std_{round(rewards[0], 1)}", ) + + +def check_reward_nonzero_std_with_fallback(args, samples: list[Sample], **kwargs): + """Prefer non-zero-std groups without triggering another sampling round.""" + + output = check_reward_nonzero_std(args, samples, **kwargs) + output.keep_when_insufficient = True + return output diff --git a/vime/rollout/fully_async_rollout.py b/vime/rollout/fully_async_rollout.py index 3ba24b74a..82144b3e6 100644 --- a/vime/rollout/fully_async_rollout.py +++ b/vime/rollout/fully_async_rollout.py @@ -11,8 +11,8 @@ :func:`generate_and_rm_group` which dispatches to those. Concurrency is sourced from ``args.vllm_server_concurrency`` and scaled by -the number of vLLM engines (``rollout_num_gpus // rollout_num_gpus_per_engine``) -to match the per-sample semaphore cap in :mod:`vime.rollout.vllm_rollout`. +the number of vllm engines to match the per-sample semaphore cap in +:mod:`vime.rollout.vllm_rollout`. The worker is intentionally oblivious to vime's higher-level pause / weight-update signalling (e.g. ``GenerateState.aborted``). Each in-flight @@ -34,6 +34,7 @@ from vime.rollout.vllm_rollout import GenerateState, generate_and_rm_group from vime.utils.async_utils import run +from vime.utils.http_utils import get_rollout_num_engines from vime.utils.types import Sample __all__ = [ @@ -54,9 +55,8 @@ def _get_global_worker(args, data_buffer) -> AsyncRolloutWorker: with _worker_lock: if _global_worker is None or not _global_worker.worker_thread.is_alive(): logger.info("starting fully-async rollout worker") - num_engines = max(1, args.rollout_num_gpus // args.rollout_num_gpus_per_engine) _global_worker = AsyncRolloutWorker( - args, data_buffer, concurrency=args.vllm_server_concurrency * num_engines + args, data_buffer, concurrency=args.vllm_server_concurrency * get_rollout_num_engines(args) ) _global_worker.start() return _global_worker @@ -82,7 +82,13 @@ def __init__(self, args, data_buffer, concurrency: int = 10): self.data_buffer = data_buffer self.concurrency = concurrency self.running = True - self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue(maxsize=1000) + # Unbounded on purpose: put() runs inside the event-loop thread (task + # done-callback), so a bounded queue that fills up would block the loop + # and freeze every in-flight generation. Backpressure lives in _loop() + # instead, which stops topping up while a full pool of completed groups + # is already waiting to be consumed. + self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue() + self.poll_interval = 1.0 self.worker_thread: threading.Thread | None = None self.state = GenerateState(args) @@ -98,9 +104,16 @@ def stop(self) -> None: if self.worker_thread and self.worker_thread.is_alive(): self.worker_thread.join(timeout=5) - def get_completed_groups(self) -> list[tuple[int, list[Sample]]]: + def get_completed_groups(self, limit: int | None = None) -> list[tuple[int, list[Sample]]]: + """Pop up to ``limit`` completed groups (all of them when ``None``). + + Callers that only need a fixed number of groups must pass ``limit`` — + anything popped beyond it would otherwise have to be thrown away, and + these groups are fully generated and reward-scored, with their prompts + already consumed from ``data_buffer``. + """ completed: list[tuple[int, list[Sample]]] = [] - while True: + while limit is None or len(completed) < limit: try: completed.append(self.output_queue.get_nowait()) except queue.Empty: @@ -132,8 +145,12 @@ async def _loop(self) -> None: logger.warning("fully-async task crashed: %r", e) active_tasks -= done - # Top up. - while len(active_tasks) < max_concurrent and self.running: + # Top up. The qsize gate is the queue's backpressure: once a + # full pool of completed groups is waiting, stop pulling new + # prompts until the training side drains some. + while ( + len(active_tasks) < max_concurrent and self.output_queue.qsize() < max_concurrent and self.running + ): groups = self.data_buffer.get_samples(1) if not groups: break @@ -151,10 +168,10 @@ async def _loop(self) -> None: task.add_done_callback(self._make_done_cb(gid)) active_tasks.add(task) - await asyncio.sleep(1) + await asyncio.sleep(self.poll_interval) except Exception as e: # noqa: BLE001 logger.exception("fully-async loop iteration error: %s", e) - await asyncio.sleep(1) + await asyncio.sleep(self.poll_interval) if active_tasks: logger.info( @@ -209,9 +226,10 @@ async def _generate_rollout_async(args, rollout_id: int, data_buffer) -> list[li LOG_EVERY = 30.0 while len(collected) < target: - # Pull whatever's done. + # Pull only what this rollout still needs; the surplus stays queued for + # the next rollout (that is the "queue stays warm" contract). drained = 0 - for gid, group in worker.get_completed_groups(): + for gid, group in worker.get_completed_groups(limit=target - len(collected)): collected[gid] = group drained += 1 @@ -238,7 +256,7 @@ def _key(group: list[Sample]) -> int: return int(idx) return 0 - out = sorted(collected.values(), key=_key)[:target] + out = sorted(collected.values(), key=_key) logger.info( "fully-async rollout %d: done in %.1fs, queue_left=%d", rollout_id, diff --git a/vime/rollout/on_policy_distillation.py b/vime/rollout/on_policy_distillation.py index cbf10041d..329ee1a6e 100644 --- a/vime/rollout/on_policy_distillation.py +++ b/vime/rollout/on_policy_distillation.py @@ -19,7 +19,7 @@ async def reward_func(args, sample, **kwargs): teacher_model = getattr(args, "opd_teacher_model", None) sampling_params = { "max_tokens": 1, - "temperature": 0, + "temperature": args.rollout_temperature, "prompt_logprobs": 1, "skip_special_tokens": False, } diff --git a/vime/rollout/rm_hub/__init__.py b/vime/rollout/rm_hub/__init__.py index b62ba3f48..2fb5e1e2f 100644 --- a/vime/rollout/rm_hub/__init__.py +++ b/vime/rollout/rm_hub/__init__.py @@ -53,6 +53,11 @@ async def remote_rm(args, sample: Sample, max_retries: int = 10): async def async_rm(args, sample: Sample, **kwargs): + # Per-sample custom_rm_path (from eval dataset config) takes priority + if sample.custom_rm_path: + rm_function = load_function(sample.custom_rm_path) + return await rm_function(args, sample, **kwargs) + if args.custom_rm_path is not None: rm_function = load_function(args.custom_rm_path) return await rm_function(args, sample, **kwargs) @@ -99,7 +104,18 @@ async def batched_async_rm( if args.custom_rm_path is not None: # Ensure the custom reward function is implemented in batch mode rm_function = load_function(args.custom_rm_path) - return await rm_function(args, samples, **kwargs) - tasks = [async_rm(args, sample, **kwargs) for sample in samples] - rewards = await asyncio.gather(*tasks) + rewards = await rm_function(args, samples, **kwargs) + else: + rewards = await asyncio.gather(*(async_rm(args, sample, **kwargs) for sample in samples)) + + rm_name = args.custom_rm_path or args.rm_type or "batched_async_rm" + if rewards is None or isinstance(rewards, (str, bytes, bytearray, dict)): + raise TypeError( + f"batched reward model {rm_name!r} returned {type(rewards).__name__} instead of an iterable of rewards" + ) + rewards = list(rewards) + if len(rewards) != len(samples): + raise ValueError( + f"batched reward model {rm_name!r} returned {len(rewards)} rewards for {len(samples)} samples" + ) return rewards diff --git a/vime/rollout/sample_hooks.py b/vime/rollout/sample_hooks.py new file mode 100644 index 000000000..bce870d9f --- /dev/null +++ b/vime/rollout/sample_hooks.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import inspect +from typing import Any + +from vime.utils.misc import load_function +from vime.utils.types import Sample + +_current_rollout_id: int | None = None + + +def set_current_rollout_id(rollout_id: int | None) -> None: + global _current_rollout_id + _current_rollout_id = rollout_id + + +def _accepted_kwargs(function, kwargs: dict[str, Any]) -> dict[str, Any]: + signature = inspect.signature(function) + if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()): + return kwargs + return {key: value for key, value in kwargs.items() if key in signature.parameters} + + +async def _apply_to_sample(args, sample: Sample, paths: list[str], **kwargs) -> Sample: + for path in paths: + hook = load_function(path) + result = hook(args, sample, **_accepted_kwargs(hook, kwargs)) + if inspect.isawaitable(result): + result = await result + if result is not None: + if not isinstance(result, Sample): + raise TypeError( + f"Rollout sample hook {path!r} returned {type(result).__name__}, expected Sample or None." + ) + sample = result + return sample + + +async def apply_rollout_sample_hooks(args, value, **kwargs): + """Apply configured hooks to every Sample leaf while preserving list shape.""" + + paths = getattr(args, "rollout_sample_hook_path", None) or [] + if not paths: + return value + kwargs.setdefault("rollout_id", _current_rollout_id) + if isinstance(value, Sample): + return await _apply_to_sample(args, value, paths, **kwargs) + if isinstance(value, list): + return [await apply_rollout_sample_hooks(args, item, **kwargs) for item in value] + raise TypeError(f"Rollout sample hooks expected Sample or list, got {type(value).__name__}.") diff --git a/vime/rollout/vllm_rollout.py b/vime/rollout/vllm_rollout.py index 3bf9b4cdc..4629ef0fd 100644 --- a/vime/rollout/vllm_rollout.py +++ b/vime/rollout/vllm_rollout.py @@ -15,30 +15,35 @@ import vllm_router # noqa: F401 — ensures vllm-router is importable on startup from tqdm import tqdm +from vime.backends.vllm_utils.server_control import abort_inflight_requests +from vime.observability.trace_utils import build_vllm_meta_trace_attrs, trace_function, trace_span from vime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput -from vime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter +from vime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter, should_drop_dynamic_filter_output +from vime.rollout.sample_hooks import apply_rollout_sample_hooks from vime.utils.async_utils import run from vime.utils.data import Dataset from vime.utils.eval_config import EvalDatasetConfig -from vime.utils.http_utils import get, post +from vime.utils.http_utils import get, get_rollout_num_engines, post from vime.utils.misc import SingletonMeta, load_function from vime.utils.processing_utils import ( + build_multimodal_messages, build_processor_kwargs, - encode_image_for_rollout_engine, load_processor, load_tokenizer, ) -from vime.utils.trace_utils import build_vllm_meta_trace_attrs, trace_function, trace_span from vime.utils.types import Sample from .rm_hub import async_rm, batched_async_rm -__all__ = ["generate_rollout", "get_model_url"] +__all__ = ["generate_rollout", "get_model_url", "prime_encoder"] logger = logging.getLogger(__name__) _PROCESSOR_PROMPT_KEYS = {"input_ids", "attention_mask"} +# Re-sweep interval while draining; bounds how long a late straggler can run. +_ABORT_RESWEEP_INTERVAL_S = 3.0 + def _coerce_flat_int_token_ids(ids: Any) -> list[int]: """Flatten tokenizer/processor output into ``list[int]`` for vLLM ``/inference/v1/generate``.""" @@ -98,17 +103,57 @@ def get_model_url(args: Namespace, model_name: str, endpoint: str = "/inference/ return f"http://{args.vllm_router_ip}:{args.vllm_router_port}{endpoint}" +def _get_image_urls(messages: list[dict[str, Any]]) -> list[str]: + return [ + image_url["url"] + for message in messages + if isinstance(message.get("content"), list) + for part in message["content"] + if isinstance(part, dict) and part.get("type") == "image_url" + for image_url in [part.get("image_url")] + if isinstance(image_url, dict) and isinstance(image_url.get("url"), str) + ] + + +async def prime_encoder(args: Namespace, messages: list[dict[str, Any]], *, model_name: str = "default") -> None: + """Make EC producers compute the images before their consumers generate.""" + image_urls = _get_image_urls(messages) + encoders = (getattr(args, "vllm_model_encoder_endpoints", None) or {}).get(model_name) + if encoders is None and model_name == "default": + metadata = getattr(args, "vllm_model_encoder_endpoints", None) or {} + if len(metadata) == 1: + encoders = next(iter(metadata.values())) + if not image_urls or encoders is None or not encoders[1]: + return + model, endpoints = encoders + + async def prime(index: int, image_url: str) -> None: + await post( + f"{endpoints[index % len(endpoints)].rstrip('/')}/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": image_url}}]}], + "max_tokens": 1, + "stream": False, + }, + headers={"x-request-id": str(uuid.uuid4())}, + ) + + await asyncio.gather(*(prime(index, url) for index, url in enumerate(image_urls))) + + class GenerateState(metaclass=SingletonMeta): - """The global state for the generation process.""" + """ + The global state for the generation process. + """ def __init__(self, args: Namespace) -> None: + # persistent state for the generation process self.args = args self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True) - self.semaphore = asyncio.Semaphore( - args.vllm_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine - ) + self.semaphore = asyncio.Semaphore(args.vllm_server_concurrency * get_rollout_num_engines(args)) self.sampling_params: dict[str, Any] = dict( temperature=args.rollout_temperature, top_p=args.rollout_top_p, @@ -120,7 +165,6 @@ def __init__(self, args: Namespace) -> None: no_stop_trim=True, spaces_between_special_tokens=False, ) - if getattr(args, "vllm_enable_deterministic_inference", False): sampling_seed_base = args.rollout_seed self.group_sampling_seeds = [sampling_seed_base + i for i in range(args.n_samples_per_prompt)] @@ -152,6 +196,7 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: for group in samples: self.pendings.add( asyncio.create_task( + # submit a group of samples as a single task. generate_and_rm_group( self.args, group, @@ -187,6 +232,22 @@ def _build_inference_sampling_params(sampling_params: dict[str, Any]) -> dict[st return sp +def _inference_generate_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]: + """Extract aligned token ids and log probabilities from a vLLM choice.""" + token_ids = choice.get("token_ids") + if not isinstance(token_ids, list) or not all(isinstance(token_id, int) for token_id in token_ids): + return [], [] + + logprobs = choice.get("logprobs") + content = logprobs.get("content") if isinstance(logprobs, dict) else [] + content = content or [] + log_probs = [ + float(content[index].get("logprob", 0.0)) if index < len(content) and isinstance(content[index], dict) else 0.0 + for index in range(len(token_ids)) + ] + return token_ids, log_probs + + def _mm_render_response_to_generate_body(render_data: Any, model: str) -> dict[str, Any]: """Turn ``/v1/chat/completions/render`` JSON into a ``/inference/v1/generate`` request body.""" if isinstance(render_data, dict) and isinstance(render_data.get("token_ids"), list): @@ -257,8 +318,7 @@ def _align_mm_feature_placeholders_to_tokens(generate_body: dict[str, Any], toke length = int(entry.get("length", -1)) if offset < 0 or length <= 0 or offset + length > len(render_token_ids): raise ValueError( - f"Cannot align vLLM {modality} placeholder: invalid render range " - f"offset={offset}, length={length}, render_len={len(render_token_ids)}" + f"Cannot align vLLM {modality} placeholder: invalid render range offset={offset}, length={length}, render_len={len(render_token_ids)}" ) ordered_entries.append((offset, str(modality), entry)) @@ -269,8 +329,7 @@ def _align_mm_feature_placeholders_to_tokens(generate_body: dict[str, Any], toke offset = _find_token_subsequence(token_ids, placeholder_tokens, search_start) if offset < 0: raise ValueError( - f"Cannot align vLLM {modality} placeholder from render offset={render_offset}, length={length}: " - "placeholder token slice not found in canonical token_ids" + f"Cannot align vLLM {modality} placeholder from render offset={render_offset}, length={length}: placeholder token slice not found in canonical token_ids" ) entry["offset"] = offset entry["length"] = len(placeholder_tokens) @@ -291,6 +350,8 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor) + sampling_params["max_new_tokens"] -= sample.response_length + assert ( sampling_params["max_new_tokens"] >= 0 ), f"max_new_tokens: {sampling_params['max_new_tokens']} should not be less than 0" @@ -300,25 +361,23 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A inference_sampling_params = _build_inference_sampling_params(sampling_params) - images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None + messages = build_multimodal_messages(sample.prompt, sample.multimodal_inputs) if not sample.tokens: sample.tokens = prompt_ids + # Use session_id for consistent hashing routing (vLLM router) headers = None if sample.session_id: if getattr(args, "router_policy", None) == "consistent_hash": headers = {"x-session-id": sample.session_id} - if images: - content: list[dict[str, Any]] = [{"type": "text", "text": sample.prompt}] - for image in images: - data_url = encode_image_for_rollout_engine(image) - content.append({"type": "image_url", "image_url": {"url": data_url}}) + if messages: render_payload = { "model": args.hf_checkpoint, - "messages": [{"role": "user", "content": content}], + "messages": messages, } + await prime_encoder(args, render_payload["messages"]) render_url = f"{base}/v1/chat/completions/render" with trace_span(sample, "vllm_mm_render", attrs={"model": args.hf_checkpoint}): render_data = await post(render_url, render_payload, headers=headers) @@ -335,7 +394,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A url = f"{base}/inference/v1/generate" payload = { "model": args.hf_checkpoint, - "token_ids": prompt_ids, + "token_ids": list(prompt_ids), "sampling_params": inference_sampling_params, } @@ -347,44 +406,14 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A choice = output["choices"][0] # Parse token_ids and logprobs from vLLM response - new_response_tokens = choice.get("token_ids") or [] - new_response_log_probs: list[float] = [] - lp = choice.get("logprobs") - if isinstance(lp, dict): - content_items = lp.get("content") or [] - new_response_log_probs = [ - float(item.get("logprob", 0.0)) if isinstance(item, dict) else 0.0 for item in content_items - ] - if not new_response_log_probs: - new_response_log_probs = [0.0] * len(new_response_tokens) + new_response_tokens, new_response_log_probs = _inference_generate_tokens_and_logprobs(choice) # Decode text from token_ids skip_sp = sampling_params.get("skip_special_tokens") skip_decode = True if skip_sp is None else bool(skip_sp) text = state.tokenizer.decode(new_response_tokens, skip_special_tokens=skip_decode) if new_response_tokens else "" - sample.tokens = sample.tokens + new_response_tokens - sample.response_length += len(new_response_tokens) - sample.response += text - - if sample.loss_mask is not None: - assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout - sample.loss_mask += [1] * len(new_response_tokens) - - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += new_response_log_probs - - if choice.get("routed_experts") is not None: - raw = base64.b64decode(choice["routed_experts"].encode("ascii"), validate=True) - arr = np.load(io.BytesIO(raw), allow_pickle=False) - sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)).reshape( - len(sample.tokens) - 1, - args.num_layers, - args.moe_router_topk, - ) - - # Build meta_info for update_from_meta_info + # Build meta_info from the vLLM `choices` response format. fr = choice.get("finish_reason") or "stop" if isinstance(fr, dict): finish = fr @@ -395,11 +424,44 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A else: finish = {"type": "stop"} meta: dict[str, Any] = {"finish_reason": finish} + if output.get("weight_version") is not None: + meta["weight_version"] = str(output["weight_version"]) usage = output.get("usage") if usage: meta["prompt_tokens"] = usage.get("prompt_tokens", 0) meta["completion_tokens"] = usage.get("completion_tokens", 0) - sample.update_from_meta_info(args, meta) + meta["cached_tokens"] = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) + spec_stats = output.get("request_spec_decode_stats") + if spec_stats: + meta["spec_accept_token_num"] = spec_stats.get( + "num_accepted_draft_tokens", spec_stats.get("num_accepted_tokens", 0) + ) + meta["spec_draft_token_num"] = spec_stats.get("num_draft_tokens", 0) + meta["spec_verify_ct"] = spec_stats.get("num_spec_steps", spec_stats.get("num_verify_steps", 0)) + + # MoE routing replay: vLLM ships routed_experts as a base64 .npy blob on the choice; + # decode here and route through meta_info. #183: guard on value (null when replay off). + routed_experts = choice.get("routed_experts") + if routed_experts is not None: + raw = base64.b64decode(routed_experts.encode("ascii"), validate=True) + meta["routed_experts"] = np.load(io.BytesIO(raw), allow_pickle=False) + + sampling_mask = choice.get("sampling_mask") + if sampling_mask is not None: + meta["top_p_token_ids"] = [token_id for token_ids in sampling_mask for token_id in token_ids] + offsets = [0] + for token_ids in sampling_mask: + offsets.append(offsets[-1] + len(token_ids)) + meta["top_p_token_offsets"] = offsets + + sample.append_response_tokens( + args, + tokens=new_response_tokens, + log_probs=new_response_log_probs, + trainable=True, + meta_info=meta, + text=text, + ) return sample @@ -444,6 +506,8 @@ async def generate_and_rm( else: sample = await generate(args, sample, sampling_params) + sample = await apply_rollout_sample_hooks(args, sample, evaluation=evaluation) + # for the rm that need the whole group, we will not do the rm here if args.group_rm: return sample @@ -456,7 +520,7 @@ async def generate_and_rm( samples_need_reward = [sample for sample in samples if sample.reward is None] with trace_span(samples_need_reward, "reward_model"): rewards = await batched_async_rm(args, samples_need_reward) - for sample, reward in zip(samples_need_reward, rewards, strict=False): + for sample, reward in zip(samples_need_reward, rewards, strict=True): sample.reward = reward return samples else: @@ -490,6 +554,7 @@ async def generate_and_rm_group( if state.aborted: return group + # Generate a unique session_id for each sample in the group for sample in group: if sample.session_id is None: sample.session_id = str(uuid.uuid4()) @@ -510,45 +575,48 @@ async def generate_and_rm_group( if not state.aborted and args.group_rm: with trace_span(group, "group_reward_model"): rewards = await batched_async_rm(args, group) - for sample, reward in zip(group, rewards, strict=False): + for sample, reward in zip(group, rewards, strict=True): sample.reward = reward return group async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: - aborted_samples: list[list[Sample]] = [] + aborted_samples = [] state = GenerateState(args) assert not state.aborted state.aborted = True - urls: list[str] = [] - paused_workers = False + loop = asyncio.get_running_loop() if state.pendings: base = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" - try: - response = await get(f"{base}/workers") - urls = [worker["url"] for worker in response["workers"]] - except Exception: - response = await get(f"{base}/list_workers") - urls = list(response["urls"]) - - logger.info(f"Abort request for {urls}") - pause_tasks = [post(f"{url.rstrip('/')}/pause?mode=abort", {}, max_retries=3) for url in urls] - pause_results = await asyncio.gather(*pause_tasks, return_exceptions=True) - for url, result in zip(urls, pause_results, strict=False): - if isinstance(result, Exception): - logger.warning(f"Failed to abort worker at {url}: {result}") - paused_workers = True + response = await get(f"{base}/workers") + urls = [worker["url"] for worker in response["workers"]] + # Delete-type abort: drop in-flight requests without pausing the scheduler. + await abort_inflight_requests(urls) + last_sweep = loop.time() + + # make sure all the pending tasks are finished count = 0 while state.pendings: - done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + done, state.pendings = await asyncio.wait( + state.pendings, + timeout=_ABORT_RESWEEP_INTERVAL_S, + return_when=asyncio.FIRST_COMPLETED, + ) + + # Re-sweep on a fixed interval to truncate late stragglers (e.g. a + # multi-turn turn-2 fired after the initial abort), regardless of drain. + if loop.time() - last_sweep >= _ABORT_RESWEEP_INTERVAL_S: + await abort_inflight_requests(urls) + last_sweep = loop.time() if not args.partial_rollout: continue + # for partial rollout, collect the partial samples into the data buffer for task in done: group = task.result() for sample in group: @@ -560,15 +628,6 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: if args.partial_rollout: logger.info(f"Collected {count} partial samples into the data buffer") - state.pendings = set() - if paused_workers: - logger.info("rollout: resuming workers after abort drain: %s", urls) - resume_tasks = [post(f"{url.rstrip('/')}/resume", {}, max_retries=3) for url in urls] - resume_results = await asyncio.gather(*resume_tasks, return_exceptions=True) - for url, result in zip(urls, resume_results, strict=False): - if isinstance(result, Exception): - logger.warning("Failed to resume worker at %s: %s", url, result) - return aborted_samples @@ -625,8 +684,13 @@ async def generate_rollout_async( assert len(group) == args.n_samples_per_prompt all_data.append(group) + dynamic_filter_output = call_dynamic_filter(dynamic_filter, args, group) - if not dynamic_filter_output.keep: + if should_drop_dynamic_filter_output( + dynamic_filter_output, + remaining_batch_size=state.remaining_batch_size, + target_data_size=target_data_size, + ): metric_gatherer.on_dynamic_filter_drop(reason=dynamic_filter_output.reason) state.remaining_batch_size -= 1 continue @@ -696,7 +760,28 @@ async def eval_rollout_single_dataset( global EVAL_PROMPT_DATASET - cache_key = dataset_cfg.cache_key + (args.hf_checkpoint, args.apply_chat_template) + eval_multimodal_keys = ( + dataset_cfg.multimodal_keys if dataset_cfg.multimodal_keys is not None else args.multimodal_keys + ) + eval_apply_chat_template = ( + dataset_cfg.apply_chat_template if dataset_cfg.apply_chat_template is not None else args.apply_chat_template + ) + eval_apply_chat_template_kwargs = ( + dataset_cfg.apply_chat_template_kwargs + if dataset_cfg.apply_chat_template_kwargs is not None + else args.apply_chat_template_kwargs + ) + + cache_key = dataset_cfg.cache_key + ( + args.hf_checkpoint, + eval_apply_chat_template, + json.dumps(eval_multimodal_keys, sort_keys=True) if eval_multimodal_keys is not None else None, + ( + json.dumps(eval_apply_chat_template_kwargs, sort_keys=True) + if eval_apply_chat_template_kwargs is not None + else None + ), + ) if cache_key not in EVAL_PROMPT_DATASET: tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) processor = load_processor(args.hf_checkpoint, trust_remote_code=True) @@ -707,11 +792,11 @@ async def eval_rollout_single_dataset( max_length=args.eval_max_prompt_len, prompt_key=dataset_cfg.input_key, label_key=dataset_cfg.label_key, - multimodal_keys=args.multimodal_keys, + multimodal_keys=eval_multimodal_keys, metadata_key=dataset_cfg.metadata_key, tool_key=dataset_cfg.tool_key, - apply_chat_template=args.apply_chat_template, - apply_chat_template_kwargs=args.apply_chat_template_kwargs, + apply_chat_template=eval_apply_chat_template, + apply_chat_template_kwargs=eval_apply_chat_template_kwargs, ) dataset = EVAL_PROMPT_DATASET[cache_key] @@ -720,22 +805,38 @@ async def eval_rollout_single_dataset( top_p=dataset_cfg.top_p, top_k=dataset_cfg.top_k, max_new_tokens=dataset_cfg.max_response_len, - stop=args.rollout_stop, - stop_token_ids=args.rollout_stop_token_ids, - skip_special_tokens=args.rollout_skip_special_tokens, - no_stop_trim=True, + stop=dataset_cfg.stop if dataset_cfg.stop is not None else args.rollout_stop, + stop_token_ids=( + dataset_cfg.stop_token_ids if dataset_cfg.stop_token_ids is not None else args.rollout_stop_token_ids + ), + skip_special_tokens=( + dataset_cfg.skip_special_tokens + if dataset_cfg.skip_special_tokens is not None + else args.rollout_skip_special_tokens + ), + no_stop_trim=dataset_cfg.no_stop_trim if dataset_cfg.no_stop_trim is not None else True, spaces_between_special_tokens=False, ) + if dataset_cfg.repetition_penalty is not None: + base_sampling_params["repetition_penalty"] = dataset_cfg.repetition_penalty + min_new_tokens = dataset_cfg.min_new_tokens + if min_new_tokens is None: + min_new_tokens = getattr(args, "eval_min_new_tokens", None) + if min_new_tokens is not None: + base_sampling_params["min_new_tokens"] = min_new_tokens tasks = [] + # do multiple samples for eval prompts sample_index = 0 for _i, prompt_sample in enumerate(dataset.samples): for j in range(dataset_cfg.n_samples_per_eval_prompt): + # use the same prompt for multiple samples sample = copy.deepcopy(prompt_sample) sample.index = sample_index sample_index += 1 sample.session_id = str(uuid.uuid4()) sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None)) + sample.custom_rm_path = dataset_cfg.custom_rm_path sample.generate_function_path = getattr(dataset_cfg, "custom_generate_function_path", None) sampling_params = base_sampling_params if getattr(args, "vllm_enable_deterministic_inference", False): @@ -760,9 +861,7 @@ async def eval_rollout_single_dataset( if do_print: logged_sample = sample[0] if isinstance(sample, list) else sample logger.info( - "eval_rollout_single_dataset example data: " - f"{[str(logged_sample.prompt) + logged_sample.response]} " - f"reward={logged_sample.reward}" + f"eval_rollout_single_dataset example data: {[str(logged_sample.prompt) + logged_sample.response]} reward={logged_sample.reward}" ) do_print = False if isinstance(sample, list): diff --git a/vime/rollout/vllm_streaming_rollout.py b/vime/rollout/vllm_streaming_rollout.py index f143499f4..9007b37dc 100644 --- a/vime/rollout/vllm_streaming_rollout.py +++ b/vime/rollout/vllm_streaming_rollout.py @@ -34,6 +34,7 @@ import numpy as np +from vime.observability.trace_utils import build_vllm_meta_trace_attrs, trace_span from vime.rollout.vllm_rollout import ( GenerateState, _align_mm_feature_placeholders_to_tokens, @@ -41,10 +42,10 @@ _coerce_flat_int_token_ids, _mm_render_response_to_generate_body, _prepare_prompt_ids, + prime_encoder, ) from vime.utils import http_utils -from vime.utils.processing_utils import build_processor_kwargs, encode_image_for_rollout_engine -from vime.utils.trace_utils import build_vllm_meta_trace_attrs, trace_span +from vime.utils.processing_utils import build_multimodal_messages, build_processor_kwargs from vime.utils.types import Sample __all__ = ["generate_streaming"] @@ -73,8 +74,8 @@ def _base_dataset_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: """Streaming counterpart to :func:`vime.rollout.vllm_rollout.generate`. - Writes the accumulated state from each SSE chunk onto ``sample`` so an abort - that cuts the stream still leaves a coherent partial sample behind. + Writes the cumulative state from each SSE chunk onto ``sample`` so an + abort that cuts the stream still leaves a coherent partial sample behind. """ if args.ci_test: assert isinstance(sample.prompt, str) @@ -83,28 +84,23 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d base = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" url = f"{base}/inference/v1/generate" - assert ( - sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED + assert sample.status in ( + Sample.Status.PENDING, + Sample.Status.ABORTED, ), f"Sample status is {sample.status}" prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor) base_prompt_ids = _base_dataset_prompt_ids(sample, state.tokenizer, state.processor) - # Multimodal samples use the same render-dance as the non-streaming text - # path (/v1/chat/completions/render → features), then stream the generate - # call. Streaming only changes how output is returned (SSE deltas vs one - # JSON); the image render (input prep) is identical. Built below once - # sampling params + token_ids are resolved. - images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None + messages = build_multimodal_messages(sample.prompt, sample.multimodal_inputs) params = dict(sampling_params) if len(sample.response) > 0: params["max_new_tokens"] -= len(sample.tokens) - len(base_prompt_ids) - assert params["max_new_tokens"] >= 0, ( - f"max_new_tokens: {params['max_new_tokens']} should not be less than 0 " - f"(after partial continuation adjustment; tokens={len(sample.tokens)}, base_prompt={len(base_prompt_ids)})" - ) + assert ( + params["max_new_tokens"] >= 0 + ), f"max_new_tokens: {params['max_new_tokens']} should not be less than 0 (after partial continuation adjustment; tokens={len(sample.tokens)}, base_prompt={len(base_prompt_ids)})" if params["max_new_tokens"] == 0: sample.status = Sample.Status.TRUNCATED return sample @@ -113,29 +109,19 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d if not sample.tokens: sample.tokens = prompt_ids - # vLLM ``/inference/v1/generate`` is token-only. On partial continuation, - # send the full prompt+response prefix so the engine continues from the - # current sample state (mirrors the non-streaming text path). if len(sample.response) > 0: token_ids = _coerce_flat_int_token_ids(sample.tokens) else: token_ids = prompt_ids - # Use session_id for consistent_hash routing (vime convention: x-session-id - # header + policy "consistent_hash"). See vllm_rollout.generate. headers = None if sample.session_id and getattr(args, "router_policy", None) == "consistent_hash": headers = {"x-session-id": sample.session_id} payload: dict[str, Any] - if images: - # Same render-dance as vllm_rollout.generate's MM path, then stream. - # mm placeholders live in the (stable) prompt prefix, so re-rendering and - # re-aligning to the current token_ids holds across partial continuations. - content: list[dict[str, Any]] = [{"type": "text", "text": sample.prompt}] - for image in images: - content.append({"type": "image_url", "image_url": {"url": encode_image_for_rollout_engine(image)}}) - render_payload = {"model": args.hf_checkpoint, "messages": [{"role": "user", "content": content}]} + if messages: + render_payload = {"model": args.hf_checkpoint, "messages": messages} + await prime_encoder(args, render_payload["messages"]) with trace_span(sample, "vllm_mm_render", attrs={"model": args.hf_checkpoint}): render_data = await http_utils.post(f"{base}/v1/chat/completions/render", render_payload, headers=headers) payload = _mm_render_response_to_generate_body(render_data, args.hf_checkpoint) @@ -152,6 +138,8 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d "stream": True, } + payload["stream_options"] = {"include_usage": True} + # Snapshot pre-call sample state. vLLM's SSE chunks are *deltas* within this # call; on each chunk we append the delta and rebuild the post-call view of # the sample = prior state + accumulated deltas. A mid-stream break leaves @@ -169,6 +157,9 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d call_log_probs: list[float] = [] last_choice: dict[str, Any] | None = None last_usage: dict[str, Any] | None = None + weight_version: str | None = None + request_spec_decode_stats: dict[str, int] | None = None + sampling_mask: list[list[int]] | None = None finish_reason: Any = None client = http_utils._http_client @@ -191,6 +182,11 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d logger.warning("vllm_streaming: skipping non-JSON chunk: %r", data_str[:120]) continue + if chunk.get("weight_version") is not None: + weight_version = str(chunk["weight_version"]) + if chunk.get("request_spec_decode_stats") is not None: + request_spec_decode_stats = chunk["request_spec_decode_stats"] + choices = chunk.get("choices") or [] if not choices: # usage-only / keepalive chunk @@ -199,14 +195,16 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d continue choice = choices[0] last_choice = choice + if choice.get("sampling_mask") is not None: + if sampling_mask is None: + sampling_mask = [] + sampling_mask.extend(choice["sampling_mask"]) if chunk.get("usage"): last_usage = chunk["usage"] if choice.get("finish_reason"): finish_reason = choice["finish_reason"] - # Each streamed choice carries only this chunk's *delta* tokens - # (GenerateResponseStreamChoice), so accumulate. Parse token_ids + - # logprobs.content inline, the same way the non-streaming generate() does. + # Each chunk carries only its delta tokens + logprobs; accumulate. delta_tokens = choice.get("token_ids") or [] delta_log_probs = [] lp = choice.get("logprobs") @@ -223,9 +221,7 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d # Surface partial state on the sample immediately. If the outer # abort path cuts us, whatever we've written so far is what - # survives. Decode the *accumulated* tokens (not the per-chunk - # delta) so multi-token characters straddling a chunk boundary - # decode correctly. + # survives. Decode accumulated (not per-chunk) tokens. sample.tokens = base_tokens + call_tokens sample.response = base_response + ( state.tokenizer.decode(call_tokens, skip_special_tokens=skip_decode) if call_tokens else "" @@ -243,8 +239,6 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d span.update(build_vllm_meta_trace_attrs({"choices": [last_choice], "usage": last_usage})) if finish_reason and last_choice is not None: - # Finalize exactly like the non-streaming path: align logprobs to tokens, - # rebuild meta + output_token_logprobs, then let Sample own status. new_response_tokens = call_tokens if len(call_log_probs) == len(call_tokens): new_response_log_probs = [float(x) for x in call_log_probs] @@ -264,27 +258,47 @@ async def generate_streaming(args: Namespace, sample: Sample, sampling_params: d else: finish = {"type": "stop"} meta: dict[str, Any] = {"finish_reason": finish} + if weight_version is not None: + meta["weight_version"] = weight_version if last_usage: meta["prompt_tokens"] = last_usage.get("prompt_tokens", 0) meta["completion_tokens"] = last_usage.get("completion_tokens", 0) + meta["cached_tokens"] = (last_usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) + if request_spec_decode_stats: + meta["spec_accept_token_num"] = request_spec_decode_stats.get( + "num_accepted_draft_tokens", request_spec_decode_stats.get("num_accepted_tokens", 0) + ) + meta["spec_draft_token_num"] = request_spec_decode_stats.get("num_draft_tokens", 0) + meta["spec_verify_ct"] = request_spec_decode_stats.get( + "num_spec_steps", request_spec_decode_stats.get("num_verify_steps", 0) + ) if new_response_tokens: meta["output_token_logprobs"] = [ [float(lp), int(tid)] for lp, tid in zip(new_response_log_probs, new_response_tokens, strict=True) ] - sample.update_from_meta_info(args, meta) - # MoE routing replay (when requested) ships on the terminal choice. Guard the - # value (not just key presence): vLLM includes ``routed_experts: null`` when - # replay is off, matching vllm_rollout.generate's #183 fix. + # MoE routing replay ships on the terminal choice as a base64 .npy blob; decode + # into meta_info. Guard on value: vLLM emits ``routed_experts: null`` when off. if last_choice.get("routed_experts") is not None: raw = base64.b64decode(last_choice["routed_experts"].encode("ascii"), validate=True) - arr = np.load(io.BytesIO(raw), allow_pickle=False) - sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)).reshape( - len(sample.tokens) - 1, - args.num_layers, - args.moe_router_topk, + meta["routed_experts"] = np.load(io.BytesIO(raw), allow_pickle=False) + if sampling_mask is not None: + top_p_meta = {"top_p_token_ids": [token_id for token_ids in sampling_mask for token_id in token_ids]} + offsets = [0] + for token_ids in sampling_mask: + offsets.append(offsets[-1] + len(token_ids)) + top_p_meta["top_p_token_offsets"] = offsets + sample._apply_meta_info( + args, + top_p_meta, + new_token_count=len(new_response_tokens), + update_terminal_info=False, ) + # tokens already accumulated above; finalize metadata only (no token re-append). + sample.append_response_tokens(args, meta_info=meta) elif state.aborted: + if weight_version is not None: + sample.weight_versions.append(weight_version) sample.status = Sample.Status.ABORTED return sample diff --git a/vime/utils/accelerator/__init__.py b/vime/utils/accelerator/__init__.py new file mode 100644 index 000000000..089ac7522 --- /dev/null +++ b/vime/utils/accelerator/__init__.py @@ -0,0 +1,394 @@ +"""Runtime-selectable, backend-neutral accelerator API for Vime. + +The module-level functions are compatibility shims for the historical +``vime.utils.accelerator`` API. New code can use :func:`get_accelerator` +when it needs capability inspection or dependency injection. +""" + +from __future__ import annotations + +import importlib +import logging +import os +import sys +import threading +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch + +from .base import Accelerator +from .cuda import CUDAAccelerator +from .musa import MUSAAccelerator +from .musa import is_musa_available as _is_musa_available + +logger = logging.getLogger(__name__) + +_MUSA_PATCH_IMPORTED = False +_MUSA_BOOTSTRAP_CHECKED = False +_ACCELERATOR: Accelerator | None = None +_SELECTION_LOCK = threading.RLock() + + +@dataclass(frozen=True) +class _BackendRegistration: + factory: Callable[[], Accelerator] + is_available: Callable[[], bool] + priority: int + communication_backends: tuple[str, ...] + + +_REGISTRY: dict[str, _BackendRegistration] = {} + + +def register_accelerator( + name: str, + factory: Callable[[], Accelerator], + is_available: Callable[[], bool] | None = None, + priority: int = 0, + communication_backends: tuple[str, ...] = (), +) -> None: + """Register a lazily constructed backend without changing Vime core.""" + normalized = name.strip().lower() + if not normalized or normalized in {"auto", "none"}: + raise ValueError("Accelerator name must be a non-empty backend name") + with _SELECTION_LOCK: + _REGISTRY[normalized] = _BackendRegistration( + factory=factory, + is_available=is_available or (lambda: factory().is_available()), + priority=priority, + communication_backends=tuple(name.lower() for name in communication_backends), + ) + + +def _append_musa_patch_path() -> None: + patch_path = os.environ.get("MUSA_PATCH_PATH") + if patch_path and patch_path not in sys.path: + sys.path.append(patch_path) + + +def _import_musa_patch() -> bool: + _append_musa_patch_path() + try: + importlib.import_module("musa_patch") + except ModuleNotFoundError as exc: + if exc.name == "musa_patch": + return False + raise RuntimeError(f"musa_patch failed because dependency {exc.name!r} is missing") from exc + except Exception as exc: + raise RuntimeError(f"musa_patch initialization failed: {exc}") from exc + return True + + +def is_musa_available() -> bool: + return _is_musa_available() + + +def is_musa_environment() -> bool: + return ( + is_musa_available() + or os.environ.get("VIME_ACCELERATOR", "").lower() == "musa" + or "MUSA_VISIBLE_DEVICES" in os.environ + or bool(os.environ.get("MUSA_PATCH_PATH")) + ) + + +def _try_import_musa_patch() -> bool: + global _MUSA_PATCH_IMPORTED + if _MUSA_PATCH_IMPORTED: + return True + if not is_musa_environment(): + return False + _MUSA_PATCH_IMPORTED = _import_musa_patch() + if not _MUSA_PATCH_IMPORTED and is_musa_environment(): + logger.warning("musa_patch is not importable; continuing without it") + return _MUSA_PATCH_IMPORTED + + +def _musa_requested() -> bool: + configured = os.environ.get("VIME_ACCELERATOR", "").lower() + if configured and configured != "auto": + return configured == "musa" + return "MUSA_VISIBLE_DEVICES" in os.environ or bool(os.environ.get("MUSA_PATCH_PATH")) + + +def _bootstrap_musa_patch_if_needed() -> bool: + """Bootstrap the patch for an already chosen MUSA backend at most once.""" + global _MUSA_BOOTSTRAP_CHECKED + if _MUSA_BOOTSTRAP_CHECKED: + return _MUSA_PATCH_IMPORTED + _MUSA_BOOTSTRAP_CHECKED = True + return _try_import_musa_patch() + + +def _cuda_available() -> bool: + try: + return bool(torch.cuda.is_available() and torch.cuda.device_count() > 0) + except (ImportError, RuntimeError): + return False + + +def _register_builtin_backends() -> None: + if "cuda" not in _REGISTRY: + register_accelerator("cuda", CUDAAccelerator, _cuda_available, priority=100, communication_backends=("nccl",)) + if "musa" not in _REGISTRY: + register_accelerator( + "musa", MUSAAccelerator, is_musa_available, priority=200, communication_backends=("mccl",) + ) + + +def _requested_name() -> str | None: + value = os.environ.get("VIME_ACCELERATOR") + if value and value.lower() != "auto": + return value.strip().lower() + if _musa_requested(): + return "musa" + return None + + +def _make_selected(name: str, explicit: bool) -> Accelerator: + _register_builtin_backends() + entry = _REGISTRY.get(name) + if entry is None: + available = ", ".join(sorted(_REGISTRY)) + raise ValueError(f"Unknown accelerator {name!r}; registered backends: {available}") + if name == "musa": + # musa_patch may expose torch.musa, so bootstrap after MUSA has been + # chosen but before validating and constructing its backend. + _bootstrap_musa_patch_if_needed() + if explicit and not entry.is_available(): + if name == "musa": + detail = ( + "torch.musa is unavailable; install a MUSA-enabled PyTorch runtime and set MUSA_PATCH_PATH if required" + ) + elif name == "cuda": + detail = "torch.cuda.is_available() is false or no CUDA device is visible" + else: + detail = "the backend availability check returned false" + raise RuntimeError(f"Requested accelerator {name!r} is unavailable: {detail}") + backend = entry.factory() + if not backend.is_available(): + raise RuntimeError(f"Accelerator backend {name!r} was selected but is unavailable at runtime") + return backend + + +def get_accelerator() -> Accelerator: + global _ACCELERATOR + if _ACCELERATOR is not None: + return _ACCELERATOR + with _SELECTION_LOCK: + if _ACCELERATOR is not None: + return _ACCELERATOR + _register_builtin_backends() + requested = _requested_name() + if requested is not None: + _ACCELERATOR = _make_selected(requested, explicit=True) + logger.info("Selected accelerator %s (explicit)", _ACCELERATOR.name) + return _ACCELERATOR + + # Highest priority wins; names break priority ties deterministically. + candidates = sorted(_REGISTRY.items(), key=lambda item: (-item[1].priority, item[0])) + for name, registration in candidates: + if registration.is_available(): + _ACCELERATOR = _make_selected(name, explicit=False) + logger.info("Selected accelerator %s (auto)", _ACCELERATOR.name) + return _ACCELERATOR + registered = ", ".join(sorted(_REGISTRY)) + raise RuntimeError( + "No usable accelerator was detected. " + f"Registered backends: {registered}. " + "Set VIME_ACCELERATOR explicitly or install a supported accelerator runtime." + ) + + +def initialize_accelerator() -> Accelerator | None: + """Finalize runtime selection when a backend is requested or available. + + Explicit requests retain ``get_accelerator``'s fail-fast behavior. An + environment without accelerator hardware remains importable for CPU-only + tooling and documentation. + """ + if _ACCELERATOR is not None: + return _ACCELERATOR + with _SELECTION_LOCK: + _register_builtin_backends() + if _requested_name() is not None or any(entry.is_available() for entry in _REGISTRY.values()): + return get_accelerator() + return None + + +def set_accelerator(accelerator: Accelerator) -> None: + global _ACCELERATOR + if not isinstance(accelerator, Accelerator): + raise TypeError(f"Expected Accelerator, got {type(accelerator).__name__}") + if not accelerator.is_available(): + raise RuntimeError(f"Cannot install unavailable accelerator backend {accelerator.name!r}") + with _SELECTION_LOCK: + _ACCELERATOR = accelerator + + +def reset_accelerator() -> None: + """Reset the singleton; intended for tests and process initialization.""" + global _ACCELERATOR + with _SELECTION_LOCK: + _ACCELERATOR = None + + +def _backend() -> Accelerator: + return get_accelerator() + + +def device_type() -> str: + return _backend().device_type + + +def accelerator_module() -> Any: + return _backend().accelerator_module() + + +def device(index: int | str | torch.device | None = None) -> torch.device: + return _backend().device(index) + + +def device_name(index: int | str | torch.device | None = None) -> str: + return _backend().device_name(index) + + +def set_device(index: int | str | torch.device) -> None: + return _backend().set_device(index) + + +def current_device() -> int | str: + return _backend().current_device() + + +def device_count() -> int: + return _backend().device_count() + + +def synchronize(device_arg: int | str | torch.device | None = None) -> None: + return _backend().synchronize(device_arg) + + +def current_stream(device_arg: int | str | torch.device | None = None) -> Any: + return _backend().current_stream(device_arg) + + +def default_stream(device_arg: int | str | torch.device | None = None) -> Any: + return _backend().default_stream(device_arg) + + +def stream(stream_arg: Any): + return _backend().stream(stream_arg) + + +def new_stream(*args, **kwargs) -> Any: + stream_type = _backend().Stream + if stream_type is None: + raise NotImplementedError(f"Accelerator {_backend().name!r} does not support streams") + return stream_type(*args, **kwargs) + + +def new_event(*args, **kwargs) -> Any: + event_type = _backend().Event + if event_type is None: + raise NotImplementedError(f"Accelerator {_backend().name!r} does not support events") + return event_type(*args, **kwargs) + + +def empty_cache() -> None: + return _backend().empty_cache() + + +def ipc_collect() -> None: + return _backend().ipc_collect() + + +def set_allocator_expandable_segments() -> bool: + return _backend().set_allocator_expandable_segments() + + +def mem_get_info(device_arg: int | str | torch.device | None = None) -> tuple[int, int]: + return _backend().mem_get_info(device_arg) + + +def memory_allocated(device_arg: int | str | torch.device | None = None) -> int: + return _backend().memory_allocated(device_arg) + + +def memory_reserved(device_arg: int | str | torch.device | None = None) -> int: + return _backend().memory_reserved(device_arg) + + +def get_device_properties(device_arg: int | str | torch.device | None = None) -> Any: + return _backend().get_device_properties(device_arg) + + +def memory_module() -> Any: + return _backend().memory_module() + + +def attach_oom_observer(callback) -> bool: + return _backend().attach_oom_observer(callback) + + +def supports(capability: str) -> bool: + return _backend().supports(capability) + + +def autocast(*args, **kwargs): + return _backend().autocast(*args, **kwargs) + + +def manual_seed(seed: int) -> None: + return _backend().manual_seed(seed) + + +def manual_seed_all(seed: int) -> None: + return _backend().manual_seed_all(seed) + + +def get_rng_state(device_arg: int | str | torch.device | None = None) -> torch.Tensor: + return _backend().get_rng_state(device_arg) + + +def set_rng_state(state: torch.Tensor, device_arg: int | str | torch.device | None = None) -> None: + return _backend().set_rng_state(state, device_arg) + + +def initial_seed() -> int: + return _backend().initial_seed() + + +def distributed_device_id(index: int | str | torch.device | None = None) -> torch.device | None: + return _backend().distributed_device_id(index) + + +def post_import_torch() -> None: + return _backend().post_import_torch() + + +def is_accelerator_backend(backend: str) -> bool: + """Return whether a distributed backend belongs to a registered device accelerator.""" + _register_builtin_backends() + normalized = backend.lower() + return any( + name in normalized for registration in _REGISTRY.values() for name in registration.communication_backends + ) + + +def process_group_backend(default: str = "nccl") -> str: + return _backend().communication_backend(default) + + +def weight_update_backend(default: str = "nccl") -> str: + return _backend().weight_update_backend(default) + + +def visible_devices_env_key() -> str: + return _backend().visible_devices_env + + +def resolve_visible_device_id(physical_device_id: int | float | str) -> int: + return _backend().resolve_visible_device_id(physical_device_id) diff --git a/vime/utils/accelerator/base.py b/vime/utils/accelerator/base.py new file mode 100644 index 000000000..bfc02c294 --- /dev/null +++ b/vime/utils/accelerator/base.py @@ -0,0 +1,179 @@ +"""Small, backend-neutral accelerator contract used by Vime. + +The contract intentionally contains only operations that Vime uses in its +runtime. Vendor modules are supplied by concrete implementations and are +never imported by this module. +""" + +from __future__ import annotations + +import abc +import os +from typing import Any + +import torch + + +class Accelerator(abc.ABC): + """Common device/runtime surface exposed to Vime code.""" + + name: str + device_type: str + communication_backend_name: str + + @abc.abstractmethod + def is_available(self) -> bool: + """Return whether this backend can actually execute on this host.""" + + @abc.abstractmethod + def device(self, index: int | str | torch.device | None = None) -> torch.device: + """Return a :class:`torch.device` for a local device index.""" + + @abc.abstractmethod + def device_name(self, index: int | str | torch.device | None = None) -> str: + """Return the canonical device string used by PyTorch APIs.""" + + @abc.abstractmethod + def set_device(self, index: int | str | torch.device) -> None: + """Select the current local device.""" + + @abc.abstractmethod + def current_device(self) -> int | str: + """Return the current local device index, or ``cpu`` for CPU.""" + + @abc.abstractmethod + def device_count(self) -> int: + """Return the number of visible devices.""" + + @abc.abstractmethod + def synchronize(self, device: int | str | torch.device | None = None) -> None: + """Synchronize work on one device or the current device.""" + + @abc.abstractmethod + def current_stream(self, device: int | str | torch.device | None = None) -> Any: + """Return the current stream, or ``None`` when streams are unsupported.""" + + def default_stream(self, device: int | str | torch.device | None = None) -> Any: + """Return the default stream, or ``None`` when streams are unsupported.""" + return None + + def stream(self, stream: Any): + """Return a context manager for a stream when the backend supports it.""" + raise NotImplementedError(f"Accelerator {self.name!r} does not support streams") + + @property + def Stream(self) -> Any: + return None + + @property + def Event(self) -> Any: + return None + + @abc.abstractmethod + def empty_cache(self) -> None: + """Release allocator-held, currently unused memory.""" + + def ipc_collect(self) -> None: + """Collect inter-process allocator state when supported.""" + return None + + def set_allocator_expandable_segments(self) -> bool: + """Configure expandable allocator segments when supported.""" + return False + + @abc.abstractmethod + def mem_get_info(self, device: int | str | torch.device | None = None) -> tuple[int, int]: + """Return ``(free_bytes, total_bytes)`` for the selected device.""" + + @abc.abstractmethod + def memory_allocated(self, device: int | str | torch.device | None = None) -> int: + """Return currently allocated device memory in bytes.""" + + @abc.abstractmethod + def memory_reserved(self, device: int | str | torch.device | None = None) -> int: + """Return allocator-reserved device memory in bytes.""" + + def get_device_properties(self, device: int | str | torch.device | None = None) -> Any: + return None + + def memory_module(self) -> Any: + """Return the backend memory namespace, if it exposes one.""" + return None + + def attach_oom_observer(self, callback) -> bool: + """Attach an OOM callback; return ``False`` when unsupported.""" + return False + + def supports(self, capability: str) -> bool: + """Return whether a named optional capability is implemented.""" + return False + + def autocast(self, *args, **kwargs): + """Return an autocast context for this backend.""" + raise NotImplementedError(f"Accelerator {self.name!r} does not support autocast") + + def manual_seed(self, seed: int) -> None: + """Seed the current device generator when supported.""" + return None + + def manual_seed_all(self, seed: int) -> None: + """Seed all device generators when supported.""" + return None + + def get_rng_state(self, device: int | str | torch.device | None = None) -> torch.Tensor: + """Return the current generator state.""" + return torch.get_rng_state() + + def set_rng_state(self, state: torch.Tensor, device: int | str | torch.device | None = None) -> None: + """Restore the current generator state.""" + torch.set_rng_state(state) + + def initial_seed(self) -> int: + return int(torch.initial_seed()) + + def distributed_device_id(self, index: int | str | torch.device | None = None) -> torch.device | None: + """Return the device id accepted by ``dist.init_process_group``.""" + return self.device(index) + + def post_import_torch(self) -> None: + """Apply an optional backend hook after third-party torch imports.""" + return None + + def communication_backend(self, default: str = "nccl") -> str: + """Map a logical default backend to this accelerator's transport.""" + return self.communication_backend_name if default == "nccl" else default + + def weight_update_backend(self, default: str = "nccl") -> str: + return self.communication_backend(default) + + @property + def visible_devices_env(self) -> str: + return "CUDA_VISIBLE_DEVICES" + + def resolve_visible_device_id(self, physical_device_id: int | float | str) -> int: + """Map a physical id to a local id under this backend's visibility env.""" + raw_value = str(physical_device_id).strip() + visible = os.environ.get(self.visible_devices_env) + if not visible: + return int(float(raw_value)) + + ids = [item.strip() for item in visible.split(",") if item.strip()] + if raw_value in ids: + return ids.index(raw_value) + + try: + value = int(float(raw_value)) + except ValueError: + value = None + if value is not None and str(value) in ids: + return ids.index(str(value)) + if value is not None and 0 <= value < len(ids): + return value + raise RuntimeError( + f"Device id {raw_value} is not valid under {self.visible_devices_env}={visible}. " + f"Expected one of {ids} (physical) or 0..{len(ids) - 1} (local)." + ) + + def accelerator_module(self) -> Any: + """Return the torch backend namespace, or ``None`` for CPU.""" + return None diff --git a/vime/utils/accelerator/cuda.py b/vime/utils/accelerator/cuda.py new file mode 100644 index 000000000..12f4a9140 --- /dev/null +++ b/vime/utils/accelerator/cuda.py @@ -0,0 +1,40 @@ +"""CUDA/ROCm accelerator implementation.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .torch_accelerator import TorchAccelerator + + +class CUDAAccelerator(TorchAccelerator): + name = "cuda" + device_type = "cuda" + communication_backend_name = "nccl" + + def _module(self) -> Any: + return torch.cuda + + def is_available(self) -> bool: + return bool(torch.cuda.is_available() and torch.cuda.device_count() > 0) + + def attach_oom_observer(self, callback) -> bool: + attach = getattr(torch._C, "_cuda_attach_out_of_memory_observer", None) + if attach is None: + return False + attach(callback) + return True + + def supports(self, capability: str) -> bool: + if capability == "nvml_affinity": + return torch.version.hip is None + if capability == "bf16": + checker = getattr(torch.cuda, "is_bf16_supported", None) + return bool(checker and checker()) + if capability in {"cuda_int4_extension", "vllm_fp8_utils", "strict_fp32_logits", "triton_kernels"}: + return True + if capability == "requires_cpu_initialization": + return torch.version.hip is not None + return super().supports(capability) diff --git a/vime/utils/accelerator/musa.py b/vime/utils/accelerator/musa.py new file mode 100644 index 000000000..b59072271 --- /dev/null +++ b/vime/utils/accelerator/musa.py @@ -0,0 +1,88 @@ +"""MUSA accelerator implementation. + +Importing this module never imports ``torch_musa``. A MUSA runtime or the +optional ``musa_patch`` bootstrap may attach ``torch.musa`` before selection. +""" + +from __future__ import annotations + +import importlib +from typing import Any + +import torch + +from .torch_accelerator import TorchAccelerator + + +def musa_module() -> Any: + return getattr(torch, "musa", None) + + +def is_musa_available() -> bool: + module = musa_module() + checker = getattr(module, "is_available", None) + return bool(module is not None and checker is not None and checker()) + + +class MUSAAccelerator(TorchAccelerator): + name = "musa" + device_type = "musa" + communication_backend_name = "mccl" + + @property + def visible_devices_env(self) -> str: + return "MUSA_VISIBLE_DEVICES" + + def _module(self) -> Any: + module = musa_module() + if module is None: + raise RuntimeError("MUSA backend requires a runtime that exposes torch.musa") + return module + + def is_available(self) -> bool: + return is_musa_available() + + def weight_update_backend(self, default: str = "nccl") -> str: + return "cpu:gloo,musa:mccl" if default == "nccl" else default + + def distributed_device_id(self, index: int | str | torch.device | None = None) -> None: + return None + + def post_import_torch(self) -> None: + try: + module = importlib.import_module("musa_patch") + except ModuleNotFoundError as exc: + if exc.name == "musa_patch": + return + raise RuntimeError(f"musa_patch failed because dependency {exc.name!r} is missing") from exc + callback = getattr(module, "patch_after_import_torch", None) + if callback is not None: + callback() + + def attach_oom_observer(self, callback) -> bool: + musa_c = getattr(self._module(), "_MUSAC", None) + attach = getattr(musa_c, "_musa_attach_out_of_memory_observer", None) + if attach is None: + return False + attach(callback) + return True + + def autocast(self, *args, **kwargs): + amp = getattr(self._module(), "amp", None) + autocast = getattr(amp, "autocast", None) + if autocast is None: + raise NotImplementedError("MUSA runtime does not expose torch.musa.amp.autocast") + return autocast(*args, **kwargs) + + def supports(self, capability: str) -> bool: + if capability in {"nvml_affinity", "vllm_fp8_utils", "strict_fp32_logits"}: + return False + if capability == "requires_cpu_initialization": + return True + if capability == "amp": + amp = getattr(self._module(), "amp", None) + return callable(getattr(amp, "autocast", None)) + if capability == "bf16": + checker = getattr(self._module(), "is_bf16_supported", None) + return bool(checker and checker()) + return super().supports(capability) diff --git a/vime/utils/accelerator/torch_accelerator.py b/vime/utils/accelerator/torch_accelerator.py new file mode 100644 index 000000000..a8b588d4a --- /dev/null +++ b/vime/utils/accelerator/torch_accelerator.py @@ -0,0 +1,163 @@ +"""Shared adapter for PyTorch accelerator namespaces.""" + +from __future__ import annotations + +import logging +import os +from typing import Any + +import torch + +from .base import Accelerator + +logger = logging.getLogger(__name__) + + +class TorchAccelerator(Accelerator): + """Delegate common CUDA-like APIs to a vendor torch namespace.""" + + def _module(self) -> Any: + raise NotImplementedError + + def accelerator_module(self) -> Any: + return self._module() + + def is_available(self) -> bool: + module = self._module() + checker = getattr(module, "is_available", None) + return bool(module is not None and checker is not None and checker()) + + def device(self, index: int | str | torch.device | None = None) -> torch.device: + return torch.device(self.device_name(index)) + + def device_name(self, index: int | str | torch.device | None = None) -> str: + if isinstance(index, torch.device): + return str(index) + if isinstance(index, str): + return index if ":" in index else f"{self.device_type}:{index}" + if index is None: + index = self.current_device() + return f"{self.device_type}:{index}" + + def set_device(self, index: int | str | torch.device) -> None: + self._module().set_device(index) + + def current_device(self) -> int: + return int(self._module().current_device()) + + def device_count(self) -> int: + return int(self._module().device_count()) + + def synchronize(self, device: int | str | torch.device | None = None) -> None: + if device is None: + self._module().synchronize() + else: + self._module().synchronize(device) + + def current_stream(self, device: int | str | torch.device | None = None) -> Any: + if device is None: + return self._module().current_stream() + return self._module().current_stream(device) + + def default_stream(self, device: int | str | torch.device | None = None) -> Any: + default_stream = getattr(self._module(), "default_stream", None) + if default_stream is None: + raise NotImplementedError(f"Accelerator {self.name!r} does not expose a default stream") + if device is None: + return default_stream() + return default_stream(device) + + def stream(self, stream: Any): + stream_context = getattr(self._module(), "stream", None) + if stream_context is None: + raise NotImplementedError(f"Accelerator {self.name!r} does not expose stream contexts") + return stream_context(stream) + + @property + def Stream(self) -> Any: + return getattr(self._module(), "Stream", None) + + @property + def Event(self) -> Any: + return getattr(self._module(), "Event", None) + + def empty_cache(self) -> None: + self._module().empty_cache() + + def ipc_collect(self) -> None: + collect = getattr(self._module(), "ipc_collect", None) + if collect is not None: + collect() + + def set_allocator_expandable_segments(self) -> bool: + value = os.getenv("VIME_ENABLE_EXPANDABLE_SEGMENTS", "0") + if value not in {"0", "1"}: + raise ValueError(f"VIME_ENABLE_EXPANDABLE_SEGMENTS must be 0 or 1, got {value!r}") + if value == "0": + return False + + memory = self.memory_module() + setter = getattr(memory, "_set_allocator_settings", None) + if setter is None: + logger.warning( + "%s memory allocator settings API is unavailable; skip expandable_segments:True", + self.name.upper(), + ) + return False + setter("expandable_segments:True") + return True + + def mem_get_info(self, device: int | str | torch.device | None = None) -> tuple[int, int]: + if device is None: + device = self.current_device() + free, total = self._module().mem_get_info(device) + return int(free), int(total) + + def memory_allocated(self, device: int | str | torch.device | None = None) -> int: + return int(self._module().memory_allocated(device)) + + def memory_reserved(self, device: int | str | torch.device | None = None) -> int: + return int(self._module().memory_reserved(device)) + + def get_device_properties(self, device: int | str | torch.device | None = None) -> Any: + if device is None: + device = self.current_device() + return self._module().get_device_properties(device) + + def memory_module(self) -> Any: + return getattr(self._module(), "memory", None) + + def autocast(self, *args, **kwargs): + return torch.autocast(self.device_type, *args, **kwargs) + + def manual_seed(self, seed: int) -> None: + self._module().manual_seed(seed) + + def manual_seed_all(self, seed: int) -> None: + self._module().manual_seed_all(seed) + + def get_rng_state(self, device: int | str | torch.device | None = None) -> torch.Tensor: + if device is None: + return self._module().get_rng_state() + return self._module().get_rng_state(device) + + def set_rng_state(self, state: torch.Tensor, device: int | str | torch.device | None = None) -> None: + if device is None: + self._module().set_rng_state(state) + else: + self._module().set_rng_state(state, device) + + def initial_seed(self) -> int: + return int(self._module().initial_seed()) + + def supports(self, capability: str) -> bool: + module = self._module() + if capability == "device_memory": + return all(hasattr(module, name) for name in ("empty_cache", "mem_get_info", "memory_allocated")) + if capability == "events": + return hasattr(module, "Event") + if capability == "rng": + return all(hasattr(module, name) for name in ("get_rng_state", "set_rng_state", "manual_seed")) + if capability == "streams": + return all(hasattr(module, name) for name in ("Stream", "current_stream", "stream")) + return capability in {"amp", "fp16"} diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 0dc1145ff..4eefc41fb 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -6,13 +6,12 @@ from typing import Any import yaml -from vllm_router.launch_router import RouterArgs from vime.backends.vllm_utils.arguments import validate_args as vllm_validate_args from vime.backends.vllm_utils.arguments import vllm_parse_args -from vime.utils.common import is_npu +from vime.backends.vllm_utils.external import apply_external_engine_info_to_args +from vime.observability.logging_utils import configure_logger from vime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list -from vime.utils.logging_utils import configure_logger logger = logging.getLogger(__name__) @@ -48,8 +47,9 @@ def add_cluster_arguments(parser): default=None, help=( "Number of GPUs for inference. Note that when using --colocate, " - "i.e. the training and the inference engines are on the same gpus, this param will be ignored and will be set as " - "actor_num_gpus_per_node * actor_num_nodes." + "i.e. the training and the inference engines are on the same gpus, this param will be set as " + "actor_num_gpus_per_node * actor_num_nodes unless it is explicitly set. " + "Set it to 0 to launch routers without local VLLM engines." ), ) parser.add_argument( @@ -99,20 +99,13 @@ def add_cluster_arguments(parser): ), ) - reset_arg(parser, "--distributed-backend", type=str, default="hccl") + reset_arg(parser, "--distributed-backend", type=str, default="nccl") reset_arg(parser, "--distributed-timeout-minutes", type=int, default=10) return parser def add_train_arguments(parser): # --train-backend is parsed early in _pre_parse_mode() and merged later. - parser.add_argument( - "--qkv-format", - type=str, - choices=["thd", "bshd"], - default="thd", - help="The qkv layout for Megatron backend.", - ) parser.add_argument( "--qwen-gdn-backend", type=str, @@ -127,16 +120,128 @@ def add_train_arguments(parser): help="Extra environment variables for training process, e.g. PyTorch memory management ones.", ) parser.add_argument( - "--train-memory-margin-bytes", - type=int, - default=1024**3, - help="Add margin for train memory allocation. By default we will reserve 1GB as margin.", + "--force-fp8-ue8m0-scale", + action="store_true", + default=False, + help=( + "Quantize block-FP8 rollout weights with power-of-two FP32 scales, " + "independent of the training GPU architecture. Blackwell-only scale " + "packing remains controlled by the rollout runtime requirements." + ), + ) + # Delta weight sync. + parser.add_argument( + "--update-weight-mode", + choices=["full", "delta"], + default="full", + help=( + "Weight sync strategy. 'full' (default) broadcasts every parameter " + "every sync. 'delta' diffs each sync against a pinned-CPU snapshot of the " + "previous one and ships only the changed bytes (disk transport only)." + ), + ) + parser.add_argument( + "--update-weight-transport", + choices=["nccl", "disk"], + default="nccl", + help=( + "Carrier for weight sync. In full mode, 'nccl' broadcasts chunks and " + "'disk' writes a complete HF checkpoint under --update-weight-disk-dir " + "before engines reload it. Delta mode is 'disk' only: each host applies the " + "published deltas into its local checkpoint and reloads via update_weights_from_disk." + ), + ) + parser.add_argument( + "--release-train", + action="store_true", + default=False, + help=( + "Release Megatron training actors during rollout and recreate them before each train step. " + "Requires disk weight sync and --save for Megatron reload." + ), + ) + parser.add_argument( + "--update-weight-disk-dir", + type=str, + default=None, + help=( + "Filesystem directory for disk-backed weight sync. In --update-weight-mode=full, " + "one complete HF checkpoint directory is written per sync. In delta mode, " + "one delta directory (changed tensors only) is written per sync." + ), + ) + parser.add_argument( + "--update-weight-disk-keep-files", + action="store_true", + default=False, + help=( + "Skip cleanup of full-checkpoint directories written by " + "--update-weight-mode=full --update-weight-transport=disk." + ), + ) + parser.add_argument( + "--update-weight-delta-encoding", + choices=["xor", "overwrite"], + default="xor", + help=( + "On-disk delta encoding for --update-weight-mode=delta --update-weight-transport=disk. " + "'xor' (default): new ^ old — smallest wire and fastest, but an involution that must be " + "applied exactly once against the correct base (applying it twice reverts). 'overwrite': " + "changed positions + new absolute values — larger, but idempotent (re-applicable any " + "number of times). Both are byte-level and dtype-blind; the engine reads the choice from " + "each version's index metadata." + ), + ) + parser.add_argument( + "--update-weight-delta-checksum", + choices=["xxh3-128", "blake3", "adler32"], + default="xxh3-128", + help=( + "Per-tensor integrity checksum for disk delta apply. The checksum is not the " + "apply bottleneck (the apply is decompress + XOR bound), so this is a digest-" + "property choice, not a speed one. 'xxh3-128' (default): widest fast non-" + "cryptographic digest, negligible accidental-corruption collisions. 'blake3': " + "cryptographic digest, for untrusted storage. 'adler32': 32-bit, for interop " + "with systems that expect it. The engine reads the choice from each version's " + "index metadata." + ), + ) + parser.add_argument( + "--custom-update-weight-post-write-path", + type=str, + default=None, + help=( + "Path to a custom function called on each trainer rank after a disk weight " + "sync's files are written (full or delta), before the engines read them — to " + "publish the writes on a non-POSIX filesystem (no cross-host visibility " + "without an explicit sync). " + "Signature: ``def hook(args, version_dir: str, rollout_engines) -> None``; the hook gates itself." + ), + ) + parser.add_argument( + "--custom-update-weight-pre-read-path", + type=str, + default=None, + help=( + "Path to a custom function called on each rollout host before it reads a " + "published disk weight version. Signature: " + "``def hook(source_dir: str, target_version: int) -> None``." + ), ) parser.add_argument( - "--megatron-to-hf-mode", - choices=["raw", "bridge"], - default="raw", - help="The method to convert megatron weights to hugging face weights for vLLM.", + "--update-weight-local-checkpoint-dir", + type=str, + default=None, + help=( + "Rollout-host-local directory (NVMe) holding a full HF checkpoint kept in " + "sync by each engine's pull_weights: every host copies a published full " + "checkpoint as-is or patches published deltas in place, and the engines " + "reload from it. Required for --update-weight-mode=delta " + "--update-weight-transport=disk; optional for full disk sync (engines then " + "pull to local disk instead of reading the shared dir directly). The " + "read-side counterpart of --custom-update-weight-post-write-path is " + "--custom-update-weight-pre-read-path." + ), ) parser.add_argument( "--custom-model-provider-path", @@ -163,7 +268,7 @@ def add_train_arguments(parser): type=str, nargs="*", default=None, - help="""List of regex patterns of parameter names to TRAIN. All other parameters will be FROZEN. + help=r"""List of regex patterns of parameter names to TRAIN. All other parameters will be FROZEN. Supports Python regex syntax (re.search). Examples: @@ -183,7 +288,7 @@ def add_train_arguments(parser): type=str, nargs="*", default=None, - help="""List of regex patterns of parameter names to FREEZE. Other parameters will remain trainable. + help=r"""List of regex patterns of parameter names to FREEZE. Other parameters will remain trainable. Supports Python regex syntax (re.search). Examples: @@ -197,6 +302,17 @@ def add_train_arguments(parser): --freeze-params-name-list linear_fc1 """, ) + reset_arg( + parser, + "--freeze-indexer", + action="store_true", + default=False, + help=( + "Freeze DSA indexer parameters while leaving the rest of the model " + "trainable. This supports both the GLM plugin indexer names and " + "Megatron's upstream DSA indexer module." + ), + ) parser.add_argument( "--allgather-cp", action="store_true", @@ -213,8 +329,8 @@ def add_rollout_arguments(parser): default=None, help=( "The huggingface checkpoint of the trained model. " - "This is used to initialize vLLM and also provide the tokenizer. " - "Note that, we will always update the parameters in vLLM with that of megatron before training, " + "This is used to initialize vllm and also provide the tokenizer. " + "Note that, we will always update the parameters in vllm with that of megatron before training, " "so you only need to provide a huggingface checkpoint that has the same architecture as the model you want to train. " "It doesn't necessary need to contain the most up-to-date parameters." ), @@ -247,7 +363,7 @@ def add_rollout_arguments(parser): "--rollout-temperature", type=float, default=1.0, - help="the temperature for the inference engine during rollout.", + help="the temperature for the inference engine during rollout. Must be > 0.", ) parser.add_argument( "--rollout-top-p", type=float, default=1.0, help="the top-p for the inference engine during rollout." @@ -337,7 +453,7 @@ def add_rollout_arguments(parser): help=( "This defines the granularity of the sampling batch in the rollout function. " "When the number of available samples falls below the target, a sampling " - "operation of size over_sampling_batch_size will be triggered." + "operation of size over_sampling_batch_size will be triggered. " "Regardless of whether partial rollout is used or filters are applied, " "the sampling granularity is always determined by this value. " "If this value is None, rollout_batch_size will be used as the default over_sampling_batch_size." @@ -349,9 +465,11 @@ def add_rollout_arguments(parser): default=None, help=( "This is the filter function for dynamic sampling. " - "It should be able to judge whether the result of a prompt should be selected or not." - "We will do dynamic filter for sampling as in DAPO. e.g. not all correct or all wrong samples." - "You could use `vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std` as an example." + "It should be able to judge whether the result of a prompt should be selected or not. " + "We will do dynamic filter for sampling as in DAPO. e.g. not all correct or all wrong samples. " + "You could use `vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std` as an example. " + "To avoid another sampling round when the oversampled candidates cannot fill rollout_batch_size, " + "use `vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std_with_fallback`." ), ) @@ -384,6 +502,16 @@ def add_rollout_arguments(parser): "This should be useful if you need to implement some special rollout logic, e.g. multi-turn, function calling." ), ) + parser.add_argument( + "--rollout-sample-hook-path", + action="append", + default=[], + help=( + "Import path to a hook applied to each generated rollout Sample before reward computation. " + "May be repeated. Hooks may be sync or async and have signature " + "hook(args, sample, *, rollout_id=None, evaluation=False) -> Sample | None." + ), + ) parser.add_argument( "--custom-rollout-log-function-path", type=str, @@ -447,10 +575,15 @@ def add_rollout_arguments(parser): ), ) parser.add_argument( - "--rollout-external", - action="store_true", - default=False, - help="Use external vLLM instances instead of launching them inside the framework.", + "--rollout-data-transport", + type=str, + choices=["object-store", "nixl"], + default="object-store", + help=( + "Transport for rollout data refs sent from rollout manager to trainer. Large rollout " + "fields are tensorized on CPU before the refs are stored. Set to nixl to transfer " + "those torch tensors via Ray NIXL." + ), ) parser.add_argument( "--rollout-external-engine-addrs", @@ -534,8 +667,9 @@ def add_data_arguments(parser): default=None, help=( "The path to the prompt data. " - "Currently we only support jsonl format, and each line should contains --input-key and --label-key, " - "which will be used as the prompt and the label respectively. " + "Supported formats are JSONL and Parquet (Parquet requires pyarrow). " + "Each record should contain --input-key and --label-key, which will be used as the prompt and " + "the label respectively. " "If you want to use a custom template, you can set --apply-chat-template to true, in that case, " "the input should be the same structure as an openai message, e.g. [{'role': 'user', 'content': 'blabla'}]. " ), @@ -607,11 +741,29 @@ def add_data_arguments(parser): action="store_true", default=False, help=( - "Balance the number of tokens between data parallel ranks with `karmarkar_karp` for verl. " + "Balance estimated training FLOPs between data parallel ranks with `karmarkar_karp`. " + "Micro-batch packing still follows the configured static/dynamic batching unless " + "`--balance-by-flops` is also set. " "Note that this may allocate the different response of the same prompt into different training steps." ), ) + parser.add_argument( + "--balance-by-flops", + action="store_true", + default=False, + help=( + "Use FLOPs-based workload estimation (coeff*L + L²) for micro-batch " + "partitioning via Karmarkar-Karp instead of first-fit token packing. " + "The linear coefficient is auto-computed from model config (hidden_size, " + "ffn_hidden_size, swiglu, MoE experts). Captures the quadratic cost of " + "attention, producing more balanced micro-batches when sequence lengths " + "vary widely. This may create micro-batches whose total tokens exceed " + "--max-tokens-per-gpu and cause OOM. Also enables --balance-data. " + "Requires --use-dynamic-batch-size." + ), + ) + parser.add_argument( "--use-dynamic-batch-size", action="store_true", @@ -740,8 +892,7 @@ def add_algo_arguments(parser): help=( "Path to save the model in HuggingFace format when using Megatron backend. " "The model will be saved to `save_hf.format(rollout_id)`. " - "In raw Megatron-to-HF mode, weights are saved with the same quantization config " - "as `--hf-checkpoint`. " + "Weights are saved with the same quantization config as `--hf-checkpoint`. " ), ) reset_arg(parser, "--seed", type=int, default=1234) @@ -813,6 +964,7 @@ def add_algo_arguments(parser): choices=[ "grpo", "gspo", + "cispo", "reinforce_plus_plus", "reinforce_plus_plus_baseline", "ppo", @@ -907,6 +1059,15 @@ def add_algo_arguments(parser): "If enabled, the optimizer's history will be cleared at the end of each rollout, which can sometimes help with training stability or fulfill specific experiment requirements." ), ) + parser.add_argument( + "--use-stateless-adam", + action="store_true", + default=False, + help=( + "Whether to use a stateless Adam optimizer that does not persist the first/second moment " + "estimates across steps. Requires --optimizer adam and --no-save-optim." + ), + ) parser.add_argument( "--use-rollout-logprobs", action="store_true", @@ -993,7 +1154,7 @@ def add_on_policy_distillation_arguments(parser): default=None, help=( "Type of on-policy distillation. " - "'vllm': Teacher log-probs are obtained from external vLLM server during rollout. " + "'vllm': Teacher log-probs are obtained from external VLLM server during rollout. " "'megatron': Teacher model is loaded via --opd-teacher-load and forwarded during training." ), ) @@ -1028,14 +1189,6 @@ def add_on_policy_distillation_arguments(parser): ) return parser - def add_router_arguments(parser): - # vllm-router's full CLI surface (~30 knobs: policy, cache_threshold, - # retries, health-check, …) under `--router-*` prefix (collision-safe). - # exclude_host_port=True because vime owns `--vllm-router-ip / --vllm-router-port` - # (defined in vime/backends/vllm_utils/arguments.py:add_vllm_router_arguments). - RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True) - return parser - # wandb def add_wandb_arguments(parser): # wandb parameters @@ -1146,7 +1299,7 @@ def add_debug_arguments(parser): "a literal file and reused across every rollout_id; a path containing {rollout_id} " "loads a per-rollout file (with eval_.pt for the eval pipeline). Unlike " "--load-debug-rollout-data, this does NOT force debug_train_only / skip_vllm -- " - "vLLM servers, router, weight_update and the colocate offload/onload dance all " + "vllm servers, router, weight_update and the colocate offload/onload dance all " "stay live, which is the point (memory measurement at long context)." ), ) @@ -1161,8 +1314,9 @@ def add_debug_arguments(parser): type=str, default=None, help=( - "Save the train data to this path for debugging. " - "The file will be saved to `save_debug_train_data.format(rollout_id)`." + "Save one train-side debug file containing all DP shards. CP-sharded fields are restored " + "to a uniform full-response format first. The path may contain `{rollout_id}` and the " + "single writer's `{rank}` placeholders." ), ) parser.add_argument( @@ -1182,13 +1336,6 @@ def add_debug_arguments(parser): type=int, default=None, ) - parser.add_argument( - "--profile-target", - type=str, - choices=["train_overall", "train_actor", "train_log_probs"], - default=["train_overall"], - nargs="+", - ) parser.add_argument( "--memory-recorder", type=str, @@ -1347,6 +1494,38 @@ def add_custom_megatron_plugins_arguments(parser): type=str, default=None, ) + parser.add_argument( + "--megatron-deepgemm-forward-layers", + nargs="+", + type=int, + default=None, + help=( + "Global zero-based decoder layers whose selected TE linears use " + "the VLLM-compatible block-FP8 DeepGEMM forward." + ), + ) + parser.add_argument( + "--megatron-deepgemm-forward-modules", + nargs="+", + default=None, + help="Optional module-name suffixes to replace in the selected dense layers.", + ) + parser.add_argument( + "--megatron-deepgemm-moe-forward-layers", + nargs="+", + type=int, + default=None, + help=( + "Global zero-based MoE decoder layers whose TEGroupedMLP uses " + "the VLLM-compatible grouped DeepGEMM forward." + ), + ) + parser.add_argument( + "--megatron-deepgemm-moe-forward-modules", + nargs="+", + default=None, + help="Optional TEGroupedMLP module-name suffixes; defaults to mlp.experts.", + ) return parser def add_mtp_training_arguments(parser): @@ -1362,6 +1541,112 @@ def add_mtp_training_arguments(parser): return parser + def add_dspark_training_arguments(parser): + """Add DSpark (semi-autoregressive speculative decoding) training arguments.""" + parser.add_argument( + "--dspark-block-size", + type=int, + default=7, + help="Number of draft tokens per block (parallel prediction width).", + ) + parser.add_argument( + "--dspark-num-draft-layers", + type=int, + default=5, + help="Number of decoder layers in the DSpark draft backbone.", + ) + parser.add_argument( + "--dspark-target-layer-ids", + type=lambda value: tuple(int(layer_id) for layer_id in value.split(",")), + default=(1, 9, 17, 25, 33), + help="Comma-separated policy layer indices to capture hidden states from.", + ) + parser.add_argument( + "--dspark-markov-rank", + type=int, + default=256, + help="Markov head embedding rank (0 to disable Markov head).", + ) + parser.add_argument( + "--dspark-markov-head-type", + type=str, + default="vanilla", + choices=["vanilla", "gated", "rnn"], + help="Markov head type.", + ) + parser.add_argument( + "--dspark-num-anchors", + type=int, + default=512, + help="Number of anchor positions to sample per sequence.", + ) + parser.add_argument( + "--dspark-mask-token-id", + type=int, + default=151669, + help="Token id used for masked positions in DSpark noise embedding.", + ) + parser.add_argument( + "--dspark-ce-loss-alpha", + type=float, + default=0.1, + help="Weight for cross-entropy loss component.", + ) + parser.add_argument( + "--dspark-l1-loss-alpha", + type=float, + default=0.9, + help="Weight for L1/TV loss component.", + ) + parser.add_argument( + "--dspark-confidence-head-alpha", + type=float, + default=1.0, + help="Weight for confidence head loss component.", + ) + parser.add_argument( + "--dspark-loss-decay-gamma", + type=float, + default=4.0, + help="Position decay gamma: weight *= exp(-pos / gamma).", + ) + parser.add_argument( + "--dspark-draft-loss-weight", + type=float, + default=1.0, + help="Weight multiplying draft loss added to policy loss.", + ) + parser.add_argument( + "--dspark-disable-confidence-head", + action="store_true", + default=False, + help="Disable confidence head (only CE + L1 loss).", + ) + parser.add_argument( + "--dspark-freeze-policy", + action="store_true", + default=False, + help="Freeze policy model during DSpark training. Detach policy " + "logits so gradient only flows to the draft model. Use when " + "RL signal is weak and policy degradation prevents draft " + "convergence.", + ) + parser.add_argument( + "--dspark-intermediate-size", + type=int, + default=0, + help="DSpark MLP intermediate size. 0 = auto (hidden_size * 2.75). " + "Set to match pre-trained checkpoint (9728 for Qwen3-4B).", + ) + parser.add_argument( + "--dspark-pretrained-model", + type=str, + default=None, + help="Path to pre-trained DSpark safetensors file. If set, draft " + "model weights are loaded from this file instead of random init.", + ) + return parser + def add_ci_arguments(parser): parser.add_argument( "--ci-test", @@ -1371,6 +1656,15 @@ def add_ci_arguments(parser): "--ci-disable-kl-checker", action="store_true", ) + parser.add_argument( + "--ci-train-rollout-logprob-abs-diff-threshold", + type=float, + default=0.1, + help=( + "Upper bound asserted on train/train_rollout_logprob_abs_diff when --ci-test is set. " + "Defaults to 0.1; tighten it (e.g. 1e-6) for deterministic train/rollout alignment gates." + ), + ) parser.add_argument( "--ci-save-grad-norm", type=str, @@ -1397,12 +1691,12 @@ def add_ci_arguments(parser): parser = add_on_policy_distillation_arguments(parser) parser = add_wandb_arguments(parser) parser = add_tensorboard_arguments(parser) - parser = add_router_arguments(parser) parser = add_debug_arguments(parser) parser = add_network_arguments(parser) parser = add_reward_model_arguments(parser) parser = add_rollout_buffer_arguments(parser) parser = add_mtp_training_arguments(parser) + parser = add_dspark_training_arguments(parser) parser = add_ci_arguments(parser) parser = add_custom_megatron_plugins_arguments(parser) reset_arg( @@ -1445,13 +1739,14 @@ def parse_args(add_custom_arguments=None): skip_vllm = pre.debug_train_only or pre.load_debug_rollout_data is not None # Phase 1: Parse vllm args independently (separate parser, parse_known_args). + # Skipped when vllm servers are not needed. vllm_ns = None if not skip_vllm: vllm_ns = vllm_parse_args() # Phase 2: Parse megatron + vime args. - # Uses ignore_unknown_args=True so that --vllm-* and pre-parsed CLI flags are - # silently ignored by the megatron parser. + # Uses ignore_unknown_args=True so that --vllm-* and pre-parsed CLI flags + # are silently ignored by the megatron parser. from vime.backends.megatron_utils.arguments import megatron_parse_args from vime.backends.megatron_utils.arguments import validate_args as megatron_validate_args @@ -1471,7 +1766,7 @@ def parse_args(add_custom_arguments=None): vime_validate_args(args) - if pre.train_backend == "megatron" and not args.debug_rollout_only: + if not args.debug_rollout_only: megatron_validate_args(args) if not args.debug_train_only: @@ -1558,11 +1853,6 @@ def parse_megatron_role_args(base_args, megatron_config_path, role): return role_args -def parse_critic_args(actor_args, megatron_config_path): - """Backward-compatible wrapper for critic-specific Megatron role parsing.""" - return parse_megatron_role_args(actor_args, megatron_config_path, role="critic") - - def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: """ Build evaluation dataset configurations from either --eval-config or --eval-prompt-data. @@ -1608,6 +1898,12 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def vime_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) + args.dspark_enabled = (getattr(args, "vllm_speculative_config", None) or {}).get("method") == "dspark" + + if args.rollout_temperature <= 0: + raise ValueError( + "--rollout-temperature must be > 0; temperature 0 is greedy decoding and is not a valid RL policy." + ) if args.kl_coef != 0 or args.use_kl_loss: if not os.path.exists(args.ref_load): @@ -1651,31 +1947,27 @@ def vime_validate_args(args): if args.opd_teacher_load is not None: raise ValueError("--opd-teacher-load is set but --use-opd is not enabled. Please add --use-opd flag.") - if args.megatron_to_hf_mode == "bridge": - if ( - args.load is not None - and os.path.exists(args.load) - and os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) - ): - # If is a Megatron checkpoint, won't use bridge to load hf weight. - pass - else: - if args.load is None: - args.load = args.ref_load or args.hf_checkpoint - # If is a HF checkpoint, set start_rollout_id to 0 here. - args.start_rollout_id = 0 - else: - if ( - args.load is None - or not os.path.exists(args.load) - or not os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) - ): - args.no_load_optim = True - args.no_load_rng = True - args.finetune = True + load_is_megatron = ( + args.load is not None + and os.path.exists(args.load) + and os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) + ) + load_is_hf = ( + args.load is not None and os.path.isdir(args.load) and os.path.exists(os.path.join(args.load, "config.json")) + ) + if load_is_hf: + from vime.backends.megatron_utils.hf_to_megatron import supports_hf_weight_loading + + load_is_hf = supports_hf_weight_loading(args.load) + if not load_is_megatron: + args.no_load_optim = True + args.no_load_rng = True + args.finetune = True + if not load_is_hf: args.load = args.ref_load - if args.ref_ckpt_step is not None: - args.ckpt_step = args.ref_ckpt_step + if args.ref_ckpt_step is not None: + args.ckpt_step = args.ref_ckpt_step + if args.start_rollout_id is None: args.start_rollout_id = 0 if args.eval_interval is not None: @@ -1710,23 +2002,43 @@ def vime_validate_args(args): if args.log_probs_max_tokens_per_gpu is None: args.log_probs_max_tokens_per_gpu = args.max_tokens_per_gpu + if getattr(args, "balance_by_flops", False): + assert args.use_dynamic_batch_size, "--balance-by-flops requires --use-dynamic-batch-size" + args.balance_data = True + if args.eps_clip_high is None: args.eps_clip_high = args.eps_clip + if args.advantage_estimator == "cispo" and args.eps_clip < 1.0: + logger.warning( + "CISPO is canonically single-sided, but --eps-clip=%s keeps the lower clip bound %s active. " + "Set --eps-clip 1.0 (and tune --eps-clip-high, e.g. 4.0) for the canonical wide setting.", + args.eps_clip, + 1.0 - args.eps_clip, + ) + if args.eval_reward_key is None: args.eval_reward_key = args.reward_key if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" - args.save_debug_train_data = f"{args.dump_details}/train_data/{{rollout_id}}_{{rank}}.pt" + args.save_debug_train_data = f"{args.dump_details}/train_data/{{rollout_id}}.pt" + + if args.save_debug_train_data is not None and args.save_debug_train_data == args.save_debug_rollout_data: + raise ValueError("--save-debug-train-data must not be equal to --save-debug-rollout-data.") if args.load_debug_rollout_data is not None: logger.info( f"load_debug_rollout_data {args.load_debug_rollout_data} is set, " - "will not instantiate vLLM servers and will only run the training process." + "will not instantiate vllm servers and will only run the training process." ) args.debug_train_only = True + args.rollout_external = args.rollout_external_engine_addrs is not None + + if args.rollout_external and not args.debug_train_only: + apply_external_engine_info_to_args(args, logger=logger) + args.use_critic = args.advantage_estimator == "ppo" # Critic always uses the same GPU count as actor. args.critic_num_gpus_per_node = args.actor_num_gpus_per_node @@ -1738,7 +2050,9 @@ def vime_validate_args(args): del args.offload if args.debug_rollout_only: - if args.colocate and (not args.rollout_num_gpus): + if args.rollout_external: + pass + elif args.colocate and args.rollout_num_gpus is None: args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes if args.num_gpus_per_node != args.actor_num_gpus_per_node: logger.info( @@ -1747,44 +2061,43 @@ def vime_validate_args(args): f"{args.actor_num_gpus_per_node} (per-physical-node GPU count)." ) args.num_gpus_per_node = args.actor_num_gpus_per_node + elif args.rollout_num_gpus == 0: + args.actor_num_gpus_per_node = 0 + args.actor_num_nodes = 0 else: args.actor_num_gpus_per_node = min(8, args.rollout_num_gpus) args.actor_num_nodes = args.rollout_num_gpus // args.actor_num_gpus_per_node args.colocate = False args.offload_train = args.offload_rollout = False - if args.train_memory_margin_bytes > 0: - logger.warning("Force train_memory_margin_bytes=0 since debug_rollout_only does not support it") - args.train_memory_margin_bytes = 0 assert not (args.debug_rollout_only and args.debug_train_only), ( "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." ) - # always true on offload for colocate at the moment. + # Colocate normally offloads Megatron between rollout and train. Release-train + # destroys Megatron actors instead, so only rollout needs memory-saver offload. if args.colocate: - if args.offload_train is None: + if args.release_train: + if args.offload_train: + logger.info("Ignoring --offload-train because --release-train releases train actors instead.") + args.offload_train = False + if args.offload_rollout is False: + logger.info("Ignoring --no-offload-rollout because colocated --release-train needs rollout offload.") + args.offload_rollout = True + elif args.offload_train is None: args.offload_train = True if args.offload_rollout is None: args.offload_rollout = True - # In colocate mode the rollout engines share the actor's physical nodes, so the - # GPUs-per-physical-node equals actor_num_gpus_per_node. --num-gpus-per-node defaults - # to 8 (an 8-GPU/node assumption); on hardware with a different per-node count (e.g. - # 4x GB200/node) that default is wrong for MULTI-NODE colocate: the rollout-engine - # addr/port allocation computes node_index via num_gpus_per_node and maps every engine - # to node 0, so worker-node engines are handed the head node's IP and fail to bind - # (OSError: [Errno 99] Cannot assign requested address). Derive the real per-node count. if args.num_gpus_per_node != args.actor_num_gpus_per_node: logger.info( f"colocate: overriding num_gpus_per_node {args.num_gpus_per_node} -> " f"actor_num_gpus_per_node {args.actor_num_gpus_per_node} (per-physical-node GPU count)." ) args.num_gpus_per_node = args.actor_num_gpus_per_node - if args.rollout_num_gpus != args.actor_num_gpus_per_node * args.actor_num_nodes: - logger.info( - f"rollout_num_gpus {args.rollout_num_gpus} != actor_num_gpus_per_node {args.actor_num_gpus_per_node} " - f"* actor_num_nodes {args.actor_num_nodes}, overriding rollout_num_gpus to match actor_num_gpus_per_node * actor_num_nodes." - ) + if args.rollout_num_gpus is None: args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes + elif args.rollout_num_gpus == 0: + logger.info("rollout_num_gpus is 0 under colocate; no local VLLM engines will be launched.") if args.offload_train is None: args.offload_train = False @@ -1794,10 +2107,7 @@ def vime_validate_args(args): if args.use_critic: args.offload_train = True - # Megatron mainline uses torch_memory_saver regions for these buffers. - # On NPU, the actor already owns the outer training region; opening nested - # NPU mem-pool regions causes beginAllocateToPool() failures. - if args.offload_train and not is_npu(): + if args.offload_train: args.disable_grad_buffers_cpu_backup = True args.disable_param_buffers_cpu_backup = True @@ -1853,16 +2163,6 @@ def vime_validate_args(args): if hasattr(args, k): logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") setattr(args, k, v) - # vllm launch_server_process distinguishes "user-supplied value" from - # "argparse default" via ``args._vllm_user_provided``. YAML overrides - # bypass argparse, so we register them explicitly here — without this, - # YAML values that happen to equal the vllm-side default (e.g. - # ``vllm_gpu_memory_utilization: 0.92``) would be treated as "default" - # and silently replaced by vime's preferred value. - if isinstance(k, str) and k.startswith("vllm_"): - if not hasattr(args, "_vllm_user_provided"): - args._vllm_user_provided = set() - args._vllm_user_provided.add(k) if args.eval_max_context_len is None: logger.info( @@ -1880,11 +2180,40 @@ def vime_validate_args(args): args.rollout_max_prompt_len <= args.rollout_max_context_len - 1 ), f"args.rollout_max_prompt_len ({args.rollout_max_prompt_len}) must be smaller than args.rollout_max_context_len ({args.rollout_max_context_len}) so that there is at least one generated token to compute loss." - if args.qkv_format == "bshd": - assert args.train_backend == "megatron", "bshd format is only supported for megatron backend." - assert ( - args.use_dynamic_batch_size is False - ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." - if args.only_train_params_name_list and args.freeze_params_name_list: raise ValueError("You can only specify ONE of: --only-train-params-name-list, or --freeze-params-name-list.") + + # disk-backed sync (full or delta) writes on the trainer and reads on the engines: needs a shared dir + if args.update_weight_transport == "disk" and not args.update_weight_disk_dir: + raise ValueError( + "--update-weight-transport=disk requires --update-weight-disk-dir to point at " + "a filesystem shared between the trainer and the rollout engines." + ) + if args.release_train: + if args.use_critic: + raise ValueError("--release-train does not support critic training yet.") + if args.keep_old_actor: + raise ValueError("--release-train does not support --keep-old-actor.") + if args.save is None: + raise ValueError("--release-train requires --save so the next Megatron actor can reload.") + if args.save_interval is None: + args.save_interval = 1 + if args.update_weight_mode != "full" or args.update_weight_transport != "disk": + raise ValueError("--release-train requires --update-weight-mode=full and --update-weight-transport=disk.") + if args.update_weight_mode == "delta": + if args.update_weight_transport != "disk": + raise ValueError( + "--update-weight-mode=delta requires --update-weight-transport=disk, " + f"got {args.update_weight_transport!r}." + ) + if args.colocate: + raise ValueError( + "--update-weight-mode=delta is not supported with --colocate. Colocate transfers " + "weights via CUDA IPC (only a handle crosses processes), so the delta bookkeeping " + "(snapshot + diff + encode) is pure overhead." + ) + if not args.update_weight_local_checkpoint_dir: + raise ValueError( + "--update-weight-mode=delta requires --update-weight-local-checkpoint-dir " + "(a rollout-host-local NVMe directory)." + ) diff --git a/vime/utils/common.py b/vime/utils/common.py deleted file mode 100644 index 3505f3e23..000000000 --- a/vime/utils/common.py +++ /dev/null @@ -1,42 +0,0 @@ -import os - -import torch - - -def get_cann_python_site_packages() -> str | None: - """Return CANN Python site-packages if ``acl`` is importable from there.""" - candidates: list[str] = [] - for env_key in ("ASCEND_TOOLKIT_HOME", "ASCEND_HOME_PATH"): - base = os.environ.get(env_key) - if not base: - continue - candidates.extend( - [ - os.path.join(base, "python", "site-packages"), - os.path.normpath(os.path.join(base, "..", "python", "site-packages")), - ] - ) - candidates.append("/usr/local/Ascend/ascend-toolkit/latest/python/site-packages") - - for path in candidates: - if os.path.isdir(os.path.join(path, "acl")): - return path - return None - - -def prepend_pythonpath(env: dict[str, str], *paths: str) -> None: - existing = env.get("PYTHONPATH", os.environ.get("PYTHONPATH", "")) - existing_parts = {part for part in existing.split(os.pathsep) if part} - prefix_parts = [path for path in paths if path and path not in existing_parts] - if prefix_parts: - env["PYTHONPATH"] = os.pathsep.join([*prefix_parts, existing] if existing else prefix_parts) - - -def is_npu() -> bool: - if not hasattr(torch, "npu"): - return False - - if not torch.npu.is_available(): - raise RuntimeError("torch_npu detected, but NPU device is not available or visible.") - - return True diff --git a/vime/utils/data.py b/vime/utils/data.py index d158ea627..a8e52f753 100644 --- a/vime/utils/data.py +++ b/vime/utils/data.py @@ -13,11 +13,10 @@ except ImportError: pq = None +from vime.observability.timer import Timer from vime.utils.types import MultimodalTypes, Sample -from .timer import Timer - -__all__ = ["Dataset"] +__all__ = ["Dataset", "get_source"] logger = logging.getLogger(__name__) @@ -63,7 +62,14 @@ def parquet_reader(p): if row_slice is not None: logger.info("read_file path=%s applying slice row_slice=%s", path, row_slice) - reader = itertools.islice(reader, row_slice.start, row_slice.stop, row_slice.step) + if (row_slice.start or 0) < 0 or (row_slice.stop or 0) < 0: + # islice forbids negative indices, but the @[...] syntax accepts + # them (e.g. "@[-100:]" = the last 100 rows). Resolving a negative + # bound needs the total row count, so materialize for this case and + # keep the streaming islice for plain non-negative slices. + reader = iter(list(reader)[row_slice]) + else: + reader = itertools.islice(reader, row_slice.start, row_slice.stop, row_slice.step) yield from reader @@ -92,27 +98,33 @@ def filter_long_prompt(origin_samples: list[Sample], tokenizer, processor, max_l # Use processor only for samples with actual multimodal content; use batched tokenizer for text-only. text_only = [] multimodal = [] - for sample in origin_samples: + for position, sample in enumerate(origin_samples): if sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()): - multimodal.append(sample) + multimodal.append((position, sample)) else: - text_only.append(sample) - filtered_samples = [] + text_only.append((position, sample)) + kept = [] if text_only: - prompts = [s.prompt for s in text_only] + prompts = [s.prompt for _, s in text_only] input_ids_list = tokenizer(prompts, add_special_tokens=False)["input_ids"] - for sample, input_ids in zip(text_only, input_ids_list, strict=True): + for (position, sample), input_ids in zip(text_only, input_ids_list, strict=True): if len(input_ids) <= max_length: - filtered_samples.append(sample) + kept.append((position, sample)) if multimodal: - from vime.utils.processing_utils import process_vision_info + from vime.utils.processing_utils import build_processor_kwargs - for sample in multimodal: - multimodal_inputs = process_vision_info(sample.prompt, processor) - processor_output = processor(text=sample.prompt, **multimodal_inputs) + for position, sample in multimodal: + processor_kwargs = build_processor_kwargs(sample.multimodal_inputs) + processor_output = processor(text=sample.prompt, **processor_kwargs) input_ids = processor_output["input_ids"][0] if len(input_ids) <= max_length: - filtered_samples.append(sample) + kept.append((position, sample)) + # The two groups are scored separately for throughput, so restore the + # dataset order here: without --rollout-shuffle the samples are consumed + # in this order, and training should not depend on which of them happen + # to carry multimodal content. + kept.sort(key=lambda position_and_sample: position_and_sample[0]) + filtered_samples = [sample for _, sample in kept] else: prompts = [sample.prompt for sample in origin_samples] input_ids_list = tokenizer(prompts, add_special_tokens=False)["input_ids"] @@ -161,7 +173,14 @@ def _build_messages(data: dict, prompt_key: str, as_conversation: bool, multimod f"Not enough {mt.name} data: more '{mt.placeholder}' placeholders in prompt " f"than {mt.name}s provided in data" ) - content_list.append({"type": mt.name, mt.name: content.pop(0)}) + item = content.pop(0) + # Support rich image config from https://github.com/QwenLM/Qwen3-VL/blob/main/README.md + # "images": [{"type": "image", "image": "path/to/img/01.jpeg", "max_pixels": 50176, "min_pixels": 50176}, {...}] + if isinstance(item, dict): + content_list.append(item) + # "images": ["path/to/img/01.jpeg", "url", "base64enc"] + else: + content_list.append({"type": mt.name, mt.name: item}) else: content_list.append({"type": "text", "text": segment}) message["content"] = content_list @@ -282,15 +301,38 @@ def __len__(self): return len(self.samples) -def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): +def process_rollout_data(rollout_data_ref, dp_rank, dp_size): assert len(rollout_data_ref) == dp_size rollout_data = ray.get(rollout_data_ref[dp_rank].inner) - partition = rollout_data.pop("partition") + # Keep `partition` in rollout_data: each local sample's position in the + # flattened rollout batch (== its index in the rollout debug dump's + # `samples`). It's a small list of ints, only read by the train debug dump / + # log-prob capture, and ignored by training (the data iterator only fetches + # requested keys), so there's no need to drop it. + partition = rollout_data["partition"] total_lengths = rollout_data["total_lengths"] # save the seqlen of the whole rollout batch Timer().seq_lens = total_lengths rollout_data["total_lengths"] = [total_lengths[i] for i in partition] + # `raw_reward` is shipped whole on purpose: log_passrate reshapes it into + # [rollout_batch_size, n_samples_per_prompt] groups, which only works on the + # full rollout batch. Metrics that pair a reward with this rank's per-sample + # lists (response_lengths, loss_masks, log_probs, ...) need the DP-local + # view instead, otherwise sample i's reward is matched against another + # sample's data. + if "raw_reward" in rollout_data: + rollout_data["local_raw_reward"] = [rollout_data["raw_reward"][i] for i in partition] + return rollout_data + + +def get_source(sample: Sample) -> str: + metadata = getattr(sample, "metadata", None) or {} + if getattr(sample, "source", None): + return sample.source + if metadata.get("source_name"): + return metadata["source_name"] + return "unknown" diff --git a/vime/utils/disk_delta.py b/vime/utils/disk_delta.py new file mode 100644 index 000000000..c3abc65b1 --- /dev/null +++ b/vime/utils/disk_delta.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import glob +import json +import os +import struct +import zlib + +import numpy as np + +# The delta phases (diff, zstd, checksum) are memory-bandwidth bound and release the GIL, +# so a thread pool over tensors recovers the bandwidth one thread leaves idle. +NUM_WORKERS = min(32, (os.cpu_count() or 8)) + +# Trainer-side helpers for disk-level delta weight sync. The receive side — materializing the +# host-local checkpoint and applying published deltas in place — lives in vLLM behind its +# /pull_weights endpoint, so it runs on every host while Vime only talks to one endpoint per engine. + + +def overwrite_encode(new: np.ndarray, changed_mask: np.ndarray) -> np.ndarray: + """The 'overwrite' delta: changed-position count (u4), positions (u4 each), then new values. + Idempotent to apply, unlike xor (an involution); the trainer picks the encoding per the docs.""" + pos = np.flatnonzero(changed_mask).astype(" None: + self._value = zlib.adler32(data, self._value) + + def hexdigest(self) -> str: + return f"{self._value:08x}" + + +def _new_hasher(algorithm: str): + if algorithm == "xxh3-128": + import xxhash + + return xxhash.xxh3_128() + if algorithm == "blake3": + import blake3 + + return blake3.blake3() + if algorithm == "adler32": + return _Adler32() + raise KeyError(f"unknown checksum algorithm {algorithm!r}") + + +def checksum(algorithm: str, buf) -> str: + hasher = _new_hasher(algorithm) + hasher.update(buf) + return hasher.hexdigest() + + +def _tensor_locations(ckpt_dir: str) -> dict[str, tuple[str, int, int]]: + """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header.""" + locations: dict[str, tuple[str, int, int]] = {} + for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")): + with open(path, "rb") as f: + (header_len,) = struct.unpack(" uint8 bytes`` that seeks straight to the + tensor — for reading many tensors without rescanning every header. KeyError if absent.""" + locations = _tensor_locations(ckpt_dir) + + def read(name: str) -> np.ndarray: + path, offset, nbytes = locations[name] + with open(path, "rb") as f: + f.seek(offset) + return np.frombuffer(f.read(nbytes), dtype=np.uint8) + + return read diff --git a/vime/utils/distributed_utils.py b/vime/utils/distributed_utils.py index 7f7e57271..af97bc14d 100644 --- a/vime/utils/distributed_utils.py +++ b/vime/utils/distributed_utils.py @@ -21,7 +21,12 @@ def init_gloo_group(): """Initialize Gloo group for distributed communication.""" global GLOO_GROUP if GLOO_GROUP is None: - GLOO_GROUP = dist.new_group(backend="gloo") + # The Megatron process-group reload path monkey-patches dist.new_group so model + # parallel subgroups can be rebuilt. This canonical CPU group has a + # separate lifecycle (it synchronizes WORLD transitions), so keep it + # raw and outside that registry. + new_group = getattr(dist, "old_new_group", dist.new_group) + GLOO_GROUP = new_group(backend="gloo") return GLOO_GROUP @@ -33,6 +38,18 @@ def get_gloo_group(): return GLOO_GROUP +def set_gloo_group(group): + """Replace the cached all-ranks Gloo group. + + Destroying the default WORLD process group also destroys every subgroup + registered with torch.distributed. The Megatron reload path uses this + setter when it temporarily replaces the NCCL WORLD group with a CPU Gloo + WORLD group, and again when NCCL is restored. + """ + global GLOO_GROUP + GLOO_GROUP = group + + # Copy from pytorch to allow creating multiple main groups. # https://github.com/pytorch/pytorch/blob/main/torch/distributed/distributed_c10d.py def init_process_group( diff --git a/vime/utils/dp_schedule.py b/vime/utils/dp_schedule.py index 283f8631b..60f040362 100644 --- a/vime/utils/dp_schedule.py +++ b/vime/utils/dp_schedule.py @@ -20,15 +20,15 @@ by splitting the largest multi-sample bins (dynamic only). 4. Distribute the ``K`` mbs across ``dp_size`` ranks, ``K / dp_size`` each, with either a strided round-robin or a Karmarkar-Karp pass on - mbs token sums. + estimated mbs FLOPs. Invariants guaranteed by :func:`build_dp_schedule` (asserted by the tests): - every DP rank runs the **same** ``num_microbatches`` per training step (required for PP sync); - - every mbs (dynamic path) holds ``<= max_tokens_per_gpu * cp_size`` - tokens, with one exception — an individual sample larger than that cap - lands alone in its own mbs (and that mbs is the only one allowed to - exceed the cap); + - every mbs (dynamic path without ``balance_by_flops``) holds + ``<= max_tokens_per_gpu * cp_size`` tokens, with one exception — an + individual sample larger than that cap lands alone in its own mbs (and + that mbs is the only one allowed to exceed the cap); - the union of per-rank sample indices equals the set of samples kept after trimming trailing rollouts (every kept sample placed exactly once); @@ -42,21 +42,37 @@ import logging from typing import Any +from vime.utils.flops_utils import calculate_fwd_flops from vime.utils.seqlen_balancing import expand_bins_by_splitting, first_fit_pack, get_seqlen_balanced_partitions logger = logging.getLogger(__name__) +def _calculate_workloads(step_lengths, args): + return [calculate_fwd_flops([sl], args) for sl in step_lengths] + + def _pack_step_into_mbs( step_lengths: list[int], *, + args: Any, use_dynamic_batch_size: bool, max_per_bin: int | None, micro_batch_size: int | None, + balance_by_flops: bool = False, ) -> list[list[int]]: """Group a step's samples into mbs. Returns ``mbs[k]`` = local indices into ``step_lengths``.""" if use_dynamic_batch_size: assert max_per_bin is not None + if balance_by_flops: + total_tokens = sum(step_lengths) + num_mbs = max(1, (total_tokens + max_per_bin - 1) // max_per_bin) + if num_mbs >= len(step_lengths): + return [[i] for i in range(len(step_lengths))] + workloads = _calculate_workloads(step_lengths, args) + # FLOPs balancing does not enforce the token cap per mbs. A + # partition can exceed max_per_bin and may OOM if the cap is tight. + return get_seqlen_balanced_partitions(workloads, num_mbs, equal_size=False) return first_fit_pack(step_lengths, max_per_bin) assert micro_batch_size is not None n = len(step_lengths) @@ -141,9 +157,11 @@ def build_dp_schedule( # ``step_mbs`` indices are LOCAL into ``sample_indices``. step_mbs = _pack_step_into_mbs( step_lengths, + args=args, use_dynamic_batch_size=args.use_dynamic_batch_size, max_per_bin=max_per_bin, micro_batch_size=getattr(args, "micro_batch_size", None), + balance_by_flops=args.balance_by_flops, ) # 2. Align mbs count to a multiple of ``align_to``. @@ -170,12 +188,12 @@ def build_dp_schedule( num_mbs_per_rank = K // dp_size num_microbatches.append(num_mbs_per_rank) - # 3. Distribute mbs across ranks: KK on mbs token sums when balance_data is on, - # otherwise a strided round-robin. Both produce ``num_mbs_per_rank`` mbs per - # rank (equal_size=True is what KK needs for PP to stay synced). + # 3. Distribute mbs across ranks: KK on estimated FLOPs when rank + # workload balancing is enabled, otherwise a strided round-robin. if args.balance_data: - mbs_token_sums = [sum(step_lengths[i] for i in bin_) for bin_ in step_mbs] - rank_mbs_idx = get_seqlen_balanced_partitions(mbs_token_sums, dp_size, equal_size=True) + step_workloads = _calculate_workloads(step_lengths, args) + mbs_weights = [sum(step_workloads[i] for i in bin_) for bin_ in step_mbs] + rank_mbs_idx = get_seqlen_balanced_partitions(mbs_weights, dp_size, equal_size=True) else: rank_mbs_idx = [list(range(r, K, dp_size)) for r in range(dp_size)] diff --git a/vime/utils/eval_config.py b/vime/utils/eval_config.py index 8342c7b20..f45f2be08 100644 --- a/vime/utils/eval_config.py +++ b/vime/utils/eval_config.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import Any _MISSING = object() @@ -33,6 +33,26 @@ "default_keys": ("max_response_len",), "arg_attrs": ("eval_max_response_len", "rollout_max_response_len"), }, + "min_eval_samples": { + "dataset_keys": ("min_eval_samples",), + "default_keys": ("min_eval_samples",), + "arg_attrs": (), + }, + "stop": { + "dataset_keys": ("stop",), + "default_keys": ("stop",), + "arg_attrs": ("rollout_stop",), + }, + "stop_token_ids": { + "dataset_keys": ("stop_token_ids",), + "default_keys": ("stop_token_ids",), + "arg_attrs": ("rollout_stop_token_ids",), + }, + "min_new_tokens": { + "dataset_keys": ("min_new_tokens",), + "default_keys": ("min_new_tokens",), + "arg_attrs": ("eval_min_new_tokens",), + }, } DATASET_SAMPLE_SPECS: dict[str, dict[str, tuple[str, ...]]] = { @@ -56,6 +76,26 @@ "default_keys": ("metadata_key",), "arg_attrs": ("metadata_key",), }, + "multimodal_keys": { + "dataset_keys": ("multimodal_keys",), + "default_keys": ("multimodal_keys",), + "arg_attrs": ("multimodal_keys",), + }, + "apply_chat_template": { + "dataset_keys": ("apply_chat_template",), + "default_keys": ("apply_chat_template",), + "arg_attrs": ("apply_chat_template",), + }, + "apply_chat_template_kwargs": { + "dataset_keys": ("apply_chat_template_kwargs",), + "default_keys": ("apply_chat_template_kwargs",), + "arg_attrs": ("apply_chat_template_kwargs",), + }, + "custom_rm_path": { + "dataset_keys": ("custom_rm_path",), + "default_keys": ("custom_rm_path",), + "arg_attrs": ("eval_custom_rm_path", "custom_rm_path"), + }, } @@ -98,12 +138,16 @@ class EvalDatasetConfig: name: str path: str rm_type: str | None = None + custom_rm_path: str | None = None # Dataset-specific overrides input_key: str | None = None label_key: str | None = None tool_key: str | None = None metadata_key: str | None = None + multimodal_keys: dict[str, str] | None = None + apply_chat_template: bool | None = None + apply_chat_template_kwargs: dict[str, Any] | None = None n_samples_per_eval_prompt: int | None = None @@ -114,6 +158,9 @@ class EvalDatasetConfig: stop: list[str] | None = None stop_token_ids: list[int] | None = None min_new_tokens: int | None = None + repetition_penalty: float | None = None + skip_special_tokens: bool | None = None + no_stop_trim: bool | None = None # per-dataset custom generate function (e.g., for tool calling) custom_generate_function_path: str | None = None @@ -123,11 +170,28 @@ class EvalDatasetConfig: app_service: str | None = None eval_task_timeout: int | None = None + min_eval_samples: int | None = None + + # Early stop: terminate eval when remaining samples < eval_early_stop_remaining + # AND no new result has been received for eval_early_stop_idle_timeout seconds. + # Both must be set (non-None) for early stop to take effect. + eval_early_stop_remaining: int | None = None + eval_early_stop_idle_timeout: float | None = None + + # Inline source config (mirrors the per-source fields in the train data JSON). + # When any of these is set, eval will treat this dataset as its own "source" + # and build a Dataset-level source_config keyed by `name`. No need to set + # --source-key or have a `source` field in the eval jsonl. + message_processor: dict[str, Any] | None = None + reward_model: dict[str, Any] | None = None + remote_environment: dict[str, Any] | None = None metadata_overrides: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: self.metadata_overrides = _ensure_metadata_overrides(self.metadata_overrides) + if self.min_eval_samples is not None and self.min_eval_samples <= 0: + raise ValueError("min_eval_samples must be positive when set.") @property def cache_key(self) -> tuple[Any, ...]: @@ -204,11 +268,29 @@ def build_eval_dataset_configs( defaults: dict[str, Any], ) -> list[EvalDatasetConfig]: defaults = defaults or {} + combined_specs = {**DATASET_RUNTIME_SPECS, **DATASET_SAMPLE_SPECS} + + # A key that is neither a spec name nor an EvalDatasetConfig field would be + # silently ignored below — the same typo inside a dataset entry raises from + # the dataclass constructor, so hold `defaults` to the same standard. + valid_default_keys = {f.name for f in fields(EvalDatasetConfig)} | { + key for spec in combined_specs.values() for key in spec["default_keys"] + } + unknown_keys = set(defaults) - valid_default_keys + if unknown_keys: + raise ValueError( + f"Unknown key(s) in eval.defaults: {sorted(unknown_keys)}. " f"Valid keys: {sorted(valid_default_keys)}." + ) + datasets: list[EvalDatasetConfig] = [] for cfg in raw_config: cfg_dict = dict(cfg or {}) - combined_specs = {**DATASET_RUNTIME_SPECS, **DATASET_SAMPLE_SPECS} _apply_dataset_field_overrides(args, cfg_dict, defaults, combined_specs) + # Fields without a spec entry (rm_type, repetition_penalty, app_service, + # ...) still honor eval.defaults: dataset entry wins, default fills in. + for key, value in defaults.items(): + if key not in combined_specs: + cfg_dict.setdefault(key, value) dataset = EvalDatasetConfig(**cfg_dict) datasets.append(dataset) return datasets diff --git a/vime/utils/external_utils/command_utils.py b/vime/utils/external_utils/command_utils.py index 49b03cc48..504772da5 100644 --- a/vime/utils/external_utils/command_utils.py +++ b/vime/utils/external_utils/command_utils.py @@ -28,10 +28,9 @@ def convert_checkpoint( dir_dst: str = "/root", hf_checkpoint: str | None = None, ): - # Platforms without torch_dist conversion (e.g. NPU, verified to fail on Ascend) - # load HF weights directly via `--megatron-to-hf-mode bridge`; nothing to convert. + # Platforms without automatic conversion use native HF loading by default. if not current_platform().torch_dist_convert: - print(f"convert_checkpoint skip on {current_platform().name} (bridge load)") + print(f"convert_checkpoint skip on {current_platform().name} (native HF load)") return hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}" @@ -59,7 +58,7 @@ def convert_checkpoint( exec_command( f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && " - f"PYTHONPATH=/root/Megatron-LM " + f"PYTHONPATH={repo_base_dir}:/root/Megatron-LM:${{PYTHONPATH:-}} " f"torchrun " f"--nproc-per-node {num_gpus_per_node} " f"{multinode_args}" @@ -138,15 +137,17 @@ def execute_train( master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1") exec_command( - # vLLM renames its subprocesses (VLLM::EngineCore / Worker_TP*), so match - # the renamed children too; the [v]/[M] brackets avoid matching pkill itself. "pkill -9 -f '[v]llm serve|VLL[M]::'; " "sleep 3; " f"{'' if external_ray else 'ray stop --force; '}" f"{'' if external_ray else 'pkill -9 ray; '}" + # cannot be run in CI, o/w kill the parent script + # TODO: do we really need this kill? (or can we instead kill vime) + # "pkill -9 python; " "pkill -9 vime; " "sleep 3; " f"{'' if external_ray else 'pkill -9 ray; '}" + # "pkill -9 python; " "pkill -9 vime; " "pkill -9 redis; " "true; " @@ -166,8 +167,9 @@ def execute_train( { "env_vars": { "PYTHONPATH": "/root/Megatron-LM/", + "RAY_USE_UVLOOP": "0", "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": str(int(check_has_nvlink())), + "NCCL_NVLS_ENABLE": os.environ.get("NCCL_NVLS_ENABLE", str(int(check_has_nvlink()))), "no_proxy": f"127.0.0.1,{master_addr}", # This is needed by megatron / torch distributed in multi-node setup "MASTER_ADDR": master_addr, @@ -193,15 +195,41 @@ def execute_train( if megatron_model_type is not None else "" ) - exec_command( - f"export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && " - f"{cmd_megatron_model_source}" - f'ray job submit --address="http://127.0.0.1:8265" ' - f"--runtime-env-json='{runtime_env_json}' " - f"-- python3 {train_script} " - f"{'${MODEL_ARGS[@]}' if megatron_model_type is not None else ''} " - f"{train_args}" - ) + model_args = "${MODEL_ARGS[@]}" if megatron_model_type is not None else "" + import torch + + if torch.version.hip is not None: + # ROCm: `ray job submit` intermittently hits a "No available agent" + # race in the ROCm container. Run the train script directly against + # the ray head started above; pass the ray runtime-env as exports. + amd_env = { + "no_proxy": f"127.0.0.1,{master_addr}", + "PYTHONUNBUFFERED": "1", + "PYTHONPATH": f"{repo_base_dir}:/root/Megatron-LM/", + "RAY_USE_UVLOOP": "0", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "MASTER_ADDR": master_addr, + **extra_env_vars, + **_parse_extra_env_vars(config.extra_env_vars), + } + import shlex + + amd_exports = " ".join(f"{k}={shlex.quote(str(v))}" for k, v in amd_env.items()) + exec_command( + f"export {amd_exports} && " + f"{cmd_megatron_model_source}" + f"python3 {train_script} {model_args} {train_args}" + ) + else: + exec_command( + f"export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && " + f"{cmd_megatron_model_source}" + f'ray job submit --address="http://127.0.0.1:8265" ' + f"--runtime-env-json='{runtime_env_json}' " + f"-- python3 {train_script} " + f"{model_args} " + f"{train_args}" + ) def _parse_extra_env_vars(text: str): @@ -250,7 +278,7 @@ def create_run_id() -> str: _warned_bool_env_var_keys = set() -# copied from SGLang +# copied from VLLM def get_bool_env_var(name: str, default: str = "false") -> bool: value = os.getenv(name, default) value = value.lower() diff --git a/vime/utils/external_utils/launch.py b/vime/utils/external_utils/launch.py index a849b7476..34368ea60 100644 --- a/vime/utils/external_utils/launch.py +++ b/vime/utils/external_utils/launch.py @@ -3,38 +3,27 @@ Used only by the test and example launch utilities (`command_utils.execute_train`), not by vime core: a `Platform` describes one accelerator for the purpose of *launching a job* — how Ray advertises its devices, the device runtime env, whether torch_dist checkpoint conversion -works, unsupported features, and how to detect it. Adding one `register(Platform(...))` in the -REGISTERED PLATFORMS block reuses the resolver, launcher, and the `execute_train` seam -unchanged; but full end-to-end support for a new accelerator may still need changes in vime -core (resource selection, backends, rollout workers), which still branches on `is_npu()`. cuda -is the default, so the GPU path is unchanged. +works, and how to construct the launch command. It adapts the core ``vime.platforms`` +selection to shell commands. Imports stay stdlib-only (torch imports are lazy) so the module is unit-testable in isolation; the actual `exec_command` calls live in `command_utils`. """ import json -import os import shlex -from collections.abc import Callable from dataclasses import dataclass, field # ── Platform contract ────────────────────────────────────────────────────── -def _never() -> bool: - return False - - @dataclass(frozen=True) class Platform: name: str ray_args: str # ray-start resource flags, "{n}"-templated with the device count env: dict = field(default_factory=dict) # device runtime env (into runtime_env + raylet) - unsupported_features: frozenset = frozenset() # declarative (e.g. {"deepep"}); not enforced by the launcher yet - torch_dist_convert: bool = True # False -> load HF weights via bridge, no conversion - detect: Callable[[], bool] = _never # True on this platform's hardware (detection fallback) + torch_dist_convert: bool = True # False -> use native HF loading without automatic conversion def ray_start_args(self, num_devices: int) -> str: return self.ray_args.format(n=num_devices) @@ -49,22 +38,9 @@ def register(platform: Platform) -> None: PLATFORMS[platform.name] = platform -def registered_platforms() -> list[Platform]: - return list(PLATFORMS.values()) - - # ── Registered platforms (add a new accelerator here) ───────────────────────── -def _detect_npu() -> bool: - try: - from vime.utils.common import is_npu - - return is_npu() - except (ImportError, RuntimeError): - return False - - register(Platform(name="cuda", ray_args="--num-gpus {n}")) # default; other fields unused for cuda register( @@ -73,9 +49,7 @@ def _detect_npu() -> bool: # vime requests NPU bundles, not GPU (see ray/placement_group.py), so advertise # the custom NPU resource rather than Ray GPU capacity. ray_args="--num-gpus 0 --resources '{{\"NPU\": {n}}}'", - detect=_detect_npu, - torch_dist_convert=False, # torch_dist conversion fails on Ascend -> bridge load - unsupported_features=frozenset({"deepep", "fp8_rollout"}), + torch_dist_convert=False, # Keep HF tests unchanged; torch_dist has a separate opt-in test. env={ "PYTHONPATH": ( "/root/Megatron-LM:/root/vime:" @@ -112,15 +86,14 @@ def _detect_npu() -> bool: def current_platform() -> Platform: - """The active platform: `VIME_TEST_DEVICE` override, else the first registered platform - whose `detect()` matches. cuda is the default when none match (GPU path unchanged).""" - override = os.environ.get("VIME_TEST_DEVICE") - if override: - return PLATFORMS[override.lower()] - for platform in registered_platforms(): - if platform.detect(): - return platform - return PLATFORMS["cuda"] + """Adapt the selected core runtime platform to the launch contract.""" + from vime.platforms import current_platform as current_runtime_platform + + selected = current_runtime_platform().name + try: + return PLATFORMS[selected] + except KeyError as exc: + raise ValueError(f"No launcher adapter is registered for platform {selected!r}") from exc def launch_commands( @@ -142,6 +115,7 @@ def launch_commands( """ extra_env = extra_env or {} all_env = {**platform.env, **extra_env} + all_env["VIME_PLATFORM"] = platform.name cmds: list = [] cmds.append( "pkill -9 -f '[v]llm serve|VLL[M]::'; sleep 3; " diff --git a/vime/utils/flops_utils.py b/vime/utils/flops_utils.py index 0ff15a743..7ac90fe83 100644 --- a/vime/utils/flops_utils.py +++ b/vime/utils/flops_utils.py @@ -1,7 +1,3 @@ -def calculate_embedding_flops(seqlen, hidden_size): - return 2 * seqlen * hidden_size - - def calculate_lm_head_flops(seqlen, hidden_size, vocab_size): return 2 * seqlen * hidden_size * vocab_size diff --git a/vime/utils/health_monitor.py b/vime/utils/health_monitor.py index 52190267a..d7da45c2c 100644 --- a/vime/utils/health_monitor.py +++ b/vime/utils/health_monitor.py @@ -3,7 +3,6 @@ import ray - logger = logging.getLogger(__name__) @@ -30,7 +29,6 @@ def __init__(self, server_group, args): self._check_timeout = args.rollout_health_check_timeout self._check_first_wait = args.rollout_health_check_first_wait self._need_first_wait = True # Need to wait after each resume - self._is_checking_enabled = False # Track if health checking should be active def start(self) -> bool: """Start the health monitor thread. Called once during initialization. @@ -79,7 +77,6 @@ def stop(self) -> None: self._thread = None self._stop_event = None self._pause_event = None - self._is_checking_enabled = False def pause(self) -> None: """Pause health checking. Called when engines are offloaded.""" @@ -87,7 +84,6 @@ def pause(self) -> None: return logger.info("Pausing health monitor...") self._pause_event.set() - self._is_checking_enabled = False def resume(self) -> None: """Resume health checking. Called when engines are onloaded.""" @@ -96,11 +92,6 @@ def resume(self) -> None: logger.info("Resuming health monitor...") self._need_first_wait = True # Need to wait after each resume self._pause_event.clear() - self._is_checking_enabled = True - - def is_checking_enabled(self) -> bool: - """Return whether health checking is currently enabled (not paused).""" - return self._is_checking_enabled def _health_monitor_loop(self) -> None: assert self._stop_event is not None diff --git a/vime/utils/http_utils.py b/vime/utils/http_utils.py index 32cdbc907..40eb78509 100644 --- a/vime/utils/http_utils.py +++ b/vime/utils/http_utils.py @@ -2,7 +2,6 @@ import ipaddress import json import logging -import multiprocessing import os import random import socket @@ -127,23 +126,6 @@ def run_router(args): return 1 -def terminate_process(process: multiprocessing.Process, timeout: float = 1.0) -> None: - """Terminate a process gracefully, with forced kill as fallback. - - Args: - process: The process to terminate - timeout: Seconds to wait for graceful termination before forcing kill - """ - if not process.is_alive(): - return - - process.terminate() - process.join(timeout=timeout) - if process.is_alive(): - process.kill() - process.join() - - _http_client: httpx.AsyncClient | None = None _client_concurrency: int = 0 @@ -198,18 +180,31 @@ async def _post(client, url, payload, max_retries=60, headers=None): return output +def get_rollout_num_engines(args) -> int: + """Return the number of rollout HTTP engines behind the router.""" + if (num_engines := getattr(args, "rollout_num_engines", None)) is not None: + return int(num_engines) + + rollout_num_gpus = getattr(args, "rollout_num_gpus", None) or 0 + rollout_num_gpus_per_engine = getattr(args, "rollout_num_gpus_per_engine", None) or 1 + if rollout_num_gpus <= 0: + return 0 + return max(1, rollout_num_gpus // rollout_num_gpus_per_engine) + + def init_http_client(args): """Initialize HTTP client and optionally enable distributed POST via Ray.""" global _http_client, _client_concurrency, _distributed_post_enabled - if not args.rollout_num_gpus: + num_engines = get_rollout_num_engines(args) + if num_engines <= 0: return - _client_concurrency = args.vllm_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + _client_concurrency = args.vllm_server_concurrency * num_engines if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), timeout=httpx.Timeout(None), - trust_env=False, # internal vLLM comm only — never route through system proxy + trust_env=False, # internal VLLM comm only — never route through system proxy ) # Optionally initialize distributed POST via Ray without changing interfaces @@ -231,6 +226,8 @@ def _init_ray_distributed_post(args): import ray from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + from vime.ray.utils import add_default_ray_env_vars + # Discover alive nodes nodes = [n for n in ray.nodes() if n.get("Alive")] if not nodes: @@ -244,7 +241,7 @@ def __init__(self, concurrency: int): self._client = httpx.AsyncClient( limits=httpx.Limits(max_connections=max(1, concurrency)), timeout=httpx.Timeout(None), - trust_env=False, # internal vLLM comm only — never route through system proxy + trust_env=False, # internal VLLM comm only — never route through system proxy ) async def do_post(self, url, payload, max_retries=60, headers=None): @@ -262,6 +259,7 @@ async def do_post(self, url, payload, max_retries=60, headers=None): actor = _HttpPosterActor.options( name=None, lifetime="detached", + runtime_env={"env_vars": add_default_ray_env_vars()}, scheduling_strategy=scheduling, max_concurrency=per_actor_conc, # Use tiny CPU to schedule diff --git a/vime/utils/mask_utils.py b/vime/utils/mask_utils.py index efe5e159f..51cb43c0b 100644 --- a/vime/utils/mask_utils.py +++ b/vime/utils/mask_utils.py @@ -91,10 +91,22 @@ def gen_multi_turn_loss_mask_qwen3( prefix_message = {"role": "user", "content": "FOR CALCULATING LOSS MASK ONLY"} prefix_token_ids = self.tokenizer.apply_chat_template([prefix_message], tokenize=True, return_dict=False) - for i, message in enumerate(messages): + i = 0 + while i < len(messages): + message = messages[i] + if message["role"] == "tool": + # Qwen templates wrap consecutive tool responses in a single user turn. + group_end = i + 1 + while group_end < len(messages) and messages[group_end]["role"] == "tool": + group_end += 1 + message_group = messages[i:group_end] + else: + group_end = i + 1 + message_group = [message] + if i == 0: tailed_message_ids = self.tokenizer.apply_chat_template( - [message, prefix_message], + message_group + [prefix_message], tokenize=True, tools=tools, return_dict=False, @@ -102,7 +114,7 @@ def gen_multi_turn_loss_mask_qwen3( message_ids = tailed_message_ids[: -len(prefix_token_ids)] else: prefixed_message_ids = self.tokenizer.apply_chat_template( - [prefix_message, message], + [prefix_message] + message_group, tokenize=True, return_dict=False, ) @@ -121,6 +133,7 @@ def gen_multi_turn_loss_mask_qwen3( all_loss_masks.extend(loss_mask) all_token_ids.extend(message_ids) + i = group_end return all_token_ids, all_loss_masks diff --git a/vime/utils/megatron_bridge_utils.py b/vime/utils/megatron_bridge_utils.py deleted file mode 100644 index c87fb5b7b..000000000 --- a/vime/utils/megatron_bridge_utils.py +++ /dev/null @@ -1,54 +0,0 @@ -from contextlib import contextmanager - -try: - from megatron.core.utils import unwrap_model -except ImportError: - unwrap_model = None - - -def patch_hf_config_for_megatron_bridge(hf_config): - configs = [] - seen_config_ids = set() - - def add_config(config): - if config is None or id(config) in seen_config_ids: - return - seen_config_ids.add(id(config)) - configs.append(config) - - add_config(hf_config) - add_config(getattr(hf_config, "config", None)) - - for config in list(configs): - add_config(getattr(config, "text_config", None)) - - for config in configs: - rope_params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) - if isinstance(rope_params, dict) and "rope_theta" in rope_params and not hasattr(config, "rope_theta"): - config.rope_theta = rope_params["rope_theta"] - - return hf_config - - -def patch_auto_bridge_hf_config(bridge): - hf_pretrained = getattr(bridge, "hf_pretrained", None) - if hf_pretrained is not None: - patch_hf_config_for_megatron_bridge(hf_pretrained) - - return bridge - - -@contextmanager -def patch_megatron_model(model): - unwrapped_model = unwrap_model(model)[0] - model_config = unwrapped_model.config - attribute_was_added = False - if not hasattr(model_config, "share_embeddings_and_output_weights"): - model_config.share_embeddings_and_output_weights = unwrapped_model.share_embeddings_and_output_weights - attribute_was_added = True - - try: - yield - finally: - if attribute_was_added: - delattr(model_config, "share_embeddings_and_output_weights") diff --git a/vime/utils/memory_utils.py b/vime/utils/memory_utils.py index 3cc04d79d..91ce5e577 100644 --- a/vime/utils/memory_utils.py +++ b/vime/utils/memory_utils.py @@ -4,67 +4,36 @@ import psutil import torch import torch.distributed as dist -from vime.utils.common import is_npu + +from vime.utils import accelerator logger = logging.getLogger(__name__) def clear_memory(clear_host_memory: bool = False): - if is_npu(): - torch.npu.synchronize() - else: - torch.cuda.synchronize() + accelerator.synchronize() gc.collect() - if not is_npu(): - torch.cuda.empty_cache() - if is_npu(): - try: - torch.npu.empty_cache() - except RuntimeError: - pass + accelerator.empty_cache() if clear_host_memory: - if is_npu(): - try: - torch.npu.empty_cache() - except RuntimeError: - pass - else: - torch._C._host_emptyCache() + torch._C._host_emptyCache() def available_memory(): - if is_npu(): - device = torch.npu.current_device() - free, total = torch.npu.mem_get_info(device) - vm = psutil.virtual_memory() - return { - "gpu": str(device), - "total_GB": _byte_to_gb(total), - "free_GB": _byte_to_gb(free), - "used_GB": _byte_to_gb(total - free), - "allocated_GB": _byte_to_gb(torch.npu.memory_allocated(device)), - "reserved_GB": _byte_to_gb(torch.npu.memory_reserved(device)), - "host_total_GB": _byte_to_gb(vm.total), - "host_available_GB": _byte_to_gb(vm.available), - "host_used_GB": _byte_to_gb(vm.used), - "host_free_GB": _byte_to_gb(vm.free), - } - else: - device = torch.cuda.current_device() - free, total = torch.cuda.mem_get_info(device) - vm = psutil.virtual_memory() - return { - "gpu": str(device), - "total_GB": _byte_to_gb(total), - "free_GB": _byte_to_gb(free), - "used_GB": _byte_to_gb(total - free), - "allocated_GB": _byte_to_gb(torch.cuda.memory_allocated(device)), - "reserved_GB": _byte_to_gb(torch.cuda.memory_reserved(device)), - "host_total_GB": _byte_to_gb(vm.total), - "host_available_GB": _byte_to_gb(vm.available), - "host_used_GB": _byte_to_gb(vm.used), - "host_free_GB": _byte_to_gb(vm.free), - } + device = accelerator.current_device() + free, total = accelerator.mem_get_info(device) + vm = psutil.virtual_memory() + return { + "gpu": str(device), + "total_GB": _byte_to_gb(total), + "free_GB": _byte_to_gb(free), + "used_GB": _byte_to_gb(total - free), + "allocated_GB": _byte_to_gb(accelerator.memory_allocated(device)), + "reserved_GB": _byte_to_gb(accelerator.memory_reserved(device)), + "host_total_GB": _byte_to_gb(vm.total), + "host_available_GB": _byte_to_gb(vm.available), + "host_used_GB": _byte_to_gb(vm.used), + "host_free_GB": _byte_to_gb(vm.free), + } def _byte_to_gb(n: int): diff --git a/vime/utils/misc.py b/vime/utils/misc.py index 5b643987c..1f55b6748 100644 --- a/vime/utils/misc.py +++ b/vime/utils/misc.py @@ -1,9 +1,41 @@ import importlib import subprocess +from collections import defaultdict +from collections.abc import Iterable +from functools import cache +from typing import Any + +import torch from vime.utils.http_utils import is_port_available +def decode_int32_meta_array(meta_info: dict[str, Any], keys: str | Iterable[str]) -> torch.Tensor | None: + if isinstance(keys, str): + keys = (keys,) + for key in keys: + if key in meta_info: + value = meta_info[key] + break + else: + return None + + if value is None: + return None + if isinstance(value, str): + import pybase64 + + value = pybase64.b64decode(value.encode("ascii")) + if isinstance(value, bytes | bytearray | memoryview): + return torch.frombuffer(bytearray(value), dtype=torch.int32) + if torch.is_tensor(value): + return value.detach().to(device="cpu", dtype=torch.int32).reshape(-1) + if hasattr(value, "flags") and not value.flags.writeable: + value = value.copy() + return torch.as_tensor(value, dtype=torch.int32).reshape(-1) + + +@cache def load_function(path): """ Load a function from a module. @@ -105,13 +137,6 @@ def inner(self): return self._inner -from collections import defaultdict -from collections.abc import Callable, Iterable -from typing import Any - -import torch - - # details: https://stackoverflow.com/questions/773/how-do-i-use-itertools-groupby def group_by(iterable, key=None): """Similar to itertools.groupby, but do not require iterable to be sorted""" @@ -119,30 +144,3 @@ def group_by(iterable, key=None): for item in iterable: ret[key(item) if key is not None else item].append(item) return dict(ret) - - -def chunk_named_params_by_size(named_params: Iterable[tuple[str, torch.Tensor]], chunk_size: int): - return _chunk_by_size( - named_params, - compute_size=lambda named_weight: named_weight[1].nbytes, - chunk_size=chunk_size, - ) - - -def _chunk_by_size(objects: Iterable[Any], compute_size: Callable[[Any], int], chunk_size: int): - bucket: list[Any] = [] - bucket_size = 0 - - for obj in objects: - obj_size = compute_size(obj) - - if bucket and (bucket_size + obj_size) >= chunk_size: - yield bucket - bucket = [] - bucket_size = 0 - - bucket.append(obj) - bucket_size += obj_size - - if bucket: - yield bucket diff --git a/vime/utils/ppo_utils.py b/vime/utils/ppo_utils.py index 92ffbc328..5f07c6c21 100644 --- a/vime/utils/ppo_utils.py +++ b/vime/utils/ppo_utils.py @@ -148,54 +148,214 @@ def compute_policy_loss( return pg_losses, clipfrac -def compute_log_probs(logits: torch.Tensor, tokens: torch.Tensor, process_group: dist.ProcessGroup | None): - # TODO: when megatron is not installed, fall back to naive implementation - from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy +@torch.compile(dynamic=True) +def compute_cispo_loss( + ppo_kl: torch.Tensor, + log_probs: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float, + eps_clip_high: float, +): + """CISPO loss from MiniMax-M1 (https://arxiv.org/abs/2506.13585, Eq. 4-5): + ``-sg(clip(ratio, 1 - eps_clip, 1 + eps_clip_high)) * advantages * log_probs``. + + Unlike PPO, the IS ratio is clipped under stop-gradient and the gradient flows + through ``log_probs``, so clipped tokens still contribute gradient. The bounds + reuse the delta-from-1 convention of ``compute_policy_loss``; canonical CISPO + disables the lower bound (``eps_clip >= 1.0``). + """ + ratio = (-ppo_kl).exp() + ratio_truncated = torch.clamp(ratio, min=1.0 - eps_clip, max=1.0 + eps_clip_high) + pg_losses = -ratio_truncated.detach() * advantages * log_probs + clipfrac = (ratio_truncated != ratio).float() + return pg_losses, clipfrac - # convert to [seq_len, batch_size, vocab_size] as expected by fused_vocab_parallel_cross_entropy - logits = logits.unsqueeze(1) - tokens = tokens.unsqueeze(1) - return -fused_vocab_parallel_cross_entropy(logits, tokens, process_group) +def _maybe_all_reduce(tensor: torch.Tensor, op: dist.ReduceOp, process_group) -> None: + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=op, group=process_group) -# from https://github.com/volcengine/verl/blob/0bdf7f469854815177e73dcfe9e420836c952e6e/verl/utils/megatron/tensor_parallel.py#L99 -class _VocabParallelEntropy(torch.autograd.Function): - @staticmethod - def forward(ctx, vocab_parallel_logits: torch.Tensor, process_group: dist.ProcessGroup) -> 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=process_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=process_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=process_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 _get_vocab_parallel_rank_size(process_group) -> tuple[int, int]: + if process_group is not None and hasattr(process_group, "rank") and hasattr(process_group, "size"): + return process_group.rank(), process_group.size() + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(group=process_group), dist.get_world_size(group=process_group) + return 0, 1 + +class _VocabParallelLogProbEntropy(torch.autograd.Function): @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, None + def forward( + ctx, + vocab_parallel_logits: torch.Tensor, + target: torch.Tensor, + log_prob_keep_mask: torch.Tensor | None, + process_group, + with_entropy: bool, + with_entropy_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + with_entropy_grad = with_entropy and with_entropy_grad + vocab_parallel_logits = vocab_parallel_logits.float() + seq_len, vocab_parallel_size = vocab_parallel_logits.shape + rank, _world_size = _get_vocab_parallel_rank_size(process_group) + vocab_start_index = rank * vocab_parallel_size + vocab_end_index = vocab_start_index + vocab_parallel_size + + target_mask = (target < vocab_start_index) | (target >= vocab_end_index) + masked_target_1d = (target - vocab_start_index).clone() + masked_target_1d[target_mask] = 0 + arange_1d = torch.arange(seq_len, device=vocab_parallel_logits.device) + + def vocab_parallel_softmax( + logits: torch.Tensor, + inplace: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + logits_max = logits.max(dim=-1, keepdim=True).values + _maybe_all_reduce(logits_max, dist.ReduceOp.MAX, process_group) + # Subtract the max for numerical stability. When ``inplace`` is set, the + # caller passed a scratch buffer it owns, so overwrite it instead of + # allocating another [seq_len, vocab] tensor. + normalized_logits = logits.sub_(logits_max) if inplace else logits - logits_max + # The normalized logit at the target position is the log-prob numerator; + # gather it (a small copy) before the in-place ``exp_`` destroys it. + predicted_logits = normalized_logits.view(-1, vocab_parallel_size)[arange_1d, masked_target_1d] + # Reuse the ``normalized_logits`` storage for exp and softmax so the whole + # softmax costs a single [seq_len, vocab] buffer instead of three. + exp_logits = normalized_logits.exp_() + sum_exp_logits = exp_logits.sum(dim=-1, keepdim=True) + _maybe_all_reduce(sum_exp_logits, dist.ReduceOp.SUM, process_group) + softmax = exp_logits.div_(sum_exp_logits) + return predicted_logits, sum_exp_logits, softmax, logits_max + + entropy = vocab_parallel_logits.new_zeros((0,)) + entropy_softmax = vocab_parallel_logits.new_empty((0,)) + sum_softmax_times_logits = vocab_parallel_logits.new_empty((0,)) + + def sum_softmax_logits(softmax: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: + if softmax.is_cuda: + # Avoid materializing the full [seq_len, vocab] product buffer. + return torch.einsum("ij,ij->i", softmax, logits).unsqueeze(-1) + return (softmax * logits).sum(dim=-1, keepdim=True) + + if log_prob_keep_mask is None: + predicted_logits, log_prob_sum_exp_logits, log_prob_softmax, log_prob_logits_max = vocab_parallel_softmax( + vocab_parallel_logits + ) + if with_entropy: + entropy_softmax = log_prob_softmax + sum_softmax_times_logits = sum_softmax_logits(entropy_softmax, vocab_parallel_logits) + _maybe_all_reduce(sum_softmax_times_logits, dist.ReduceOp.SUM, process_group) + entropy = log_prob_logits_max + log_prob_sum_exp_logits.log() - sum_softmax_times_logits + entropy = entropy.squeeze(dim=-1) + else: + if with_entropy: + _entropy_predicted_logits, entropy_sum_exp_logits, entropy_softmax, entropy_logits_max = ( + vocab_parallel_softmax(vocab_parallel_logits) + ) + sum_softmax_times_logits = sum_softmax_logits(entropy_softmax, vocab_parallel_logits) + _maybe_all_reduce(sum_softmax_times_logits, dist.ReduceOp.SUM, process_group) + entropy = entropy_logits_max + entropy_sum_exp_logits.log() - sum_softmax_times_logits + entropy = entropy.squeeze(dim=-1) + + local_target_rows = torch.nonzero(~target_mask, as_tuple=False).squeeze(-1) + log_prob_logits = vocab_parallel_logits.masked_fill(~log_prob_keep_mask, float("-inf")) + if local_target_rows.numel() > 0: + log_prob_logits[local_target_rows, masked_target_1d[local_target_rows]] = vocab_parallel_logits[ + local_target_rows, masked_target_1d[local_target_rows] + ] + # ``log_prob_logits`` is an owned scratch buffer here, so let the softmax + # consume it in place rather than allocating another copy. + predicted_logits, log_prob_sum_exp_logits, log_prob_softmax, _log_prob_logits_max = vocab_parallel_softmax( + log_prob_logits, inplace=True + ) + predicted_logits = predicted_logits.masked_fill_(target_mask, 0.0).unsqueeze(-1) + _maybe_all_reduce(predicted_logits, dist.ReduceOp.SUM, process_group) + log_prob = predicted_logits - log_prob_sum_exp_logits.log() + + if not with_entropy_grad: + ctx.mark_non_differentiable(entropy) + + ctx.with_entropy_grad = with_entropy_grad + # Metric-only entropy still returns values, but does not need the + # full-vocab entropy tensors kept alive for backward. + saved_entropy_softmax = entropy_softmax if with_entropy_grad else vocab_parallel_logits.new_empty((0,)) + saved_sum_softmax_times_logits = ( + sum_softmax_times_logits if with_entropy_grad else vocab_parallel_logits.new_empty((0,)) + ) + saved_logits = vocab_parallel_logits if with_entropy_grad else vocab_parallel_logits.new_empty((0,)) + ctx.save_for_backward( + log_prob_softmax, + target_mask, + masked_target_1d, + saved_entropy_softmax, + saved_sum_softmax_times_logits, + saved_logits, + ) + return log_prob, entropy -def compute_entropy_from_logits(logits: torch.Tensor, process_group) -> torch.Tensor: - return _VocabParallelEntropy.apply(logits, process_group) + @staticmethod + def backward( + ctx, grad_log_prob: torch.Tensor | None, grad_entropy: torch.Tensor | None + ) -> tuple[torch.Tensor, None, None, None, None, None]: + ( + log_prob_softmax, + target_mask, + masked_target_1d, + entropy_softmax, + sum_softmax_times_logits, + vocab_parallel_logits, + ) = ctx.saved_tensors + + if grad_log_prob is None: + raise RuntimeError( + "_VocabParallelLogProbEntropy expected a materialized grad_log_prob. " + "Do not call ctx.set_materialize_grads(False)." + ) + + grad_entropy_input = None + if ctx.with_entropy_grad and grad_entropy is not None and grad_entropy.numel() > 0: + # In the unmasked path, entropy_softmax aliases log_prob_softmax. + # Build entropy grad before mutating log_prob_softmax below. + grad_entropy_input = sum_softmax_times_logits - vocab_parallel_logits + grad_entropy_input.mul_(entropy_softmax) + grad_entropy_input.mul_(grad_entropy.reshape(-1, 1)) + + vocab_parallel_size = log_prob_softmax.size(-1) + grad_input = log_prob_softmax.neg_() + grad_2d = grad_input.view(-1, vocab_parallel_size) + arange_1d = torch.arange(grad_2d.size(0), device=grad_2d.device) + target_update = (~target_mask).to(dtype=grad_2d.dtype) + grad_2d[arange_1d, masked_target_1d] += target_update + grad_input.mul_(grad_log_prob.reshape(-1, 1)) + + if grad_entropy_input is not None: + grad_input.add_(grad_entropy_input) + + return grad_input, None, None, None, None, None + + +def _calculate_log_probs_and_entropy_chunk( + logits: torch.Tensor, + tokens: torch.Tensor, + tp_group, + *, + with_entropy: bool, + with_entropy_grad: bool = True, + log_prob_keep_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + log_prob, entropy = _VocabParallelLogProbEntropy.apply( + logits, + tokens, + log_prob_keep_mask, + tp_group, + with_entropy, + with_entropy_grad, + ) + if not with_entropy: + entropy = None + return log_prob, entropy def get_grpo_returns( @@ -237,7 +397,7 @@ def get_reinforce_plus_plus_returns( cp_size = mpu.get_context_parallel_world_size() - final_returns_chunks = [] + token_level_rewards = [] for i in range(len(rewards)): local_kl_chunk = kl[i] total_len, response_len = total_lengths[i], response_lengths[i] @@ -254,21 +414,28 @@ def get_reinforce_plus_plus_returns( full_mask = loss_masks[i] assert full_mask.sum().item() > 0, f"Sequence at index {i} is fully masked." masked_kl = full_kl_response * full_mask - token_level_rewards = -kl_coef * masked_kl + rewards_for_seq = -kl_coef * masked_kl last_idx = full_mask.nonzero(as_tuple=True)[0][-1] - token_level_rewards[last_idx] += rewards[i] + rewards_for_seq[last_idx] += rewards[i] + token_level_rewards.append(rewards_for_seq) + + if not token_level_rewards: + return [] + + max_len = max(rewards_for_seq.size(0) for rewards_for_seq in token_level_rewards) + padded_rewards = token_level_rewards[0].new_zeros(len(token_level_rewards), max_len) + for i, rewards_for_seq in enumerate(token_level_rewards): + padded_rewards[i, : rewards_for_seq.size(0)] = rewards_for_seq - returns_for_seq = torch.zeros_like(token_level_rewards) - running_return = 0.0 - for t in reversed(range(token_level_rewards.size(0))): - # G_t = r_t + gamma * G_{t+1} - running_return = token_level_rewards[t] + gamma * running_return - returns_for_seq[t] = running_return + padded_returns = chunked_discounted_returns(padded_rewards, gamma) - # Step 4: Pick up the results corresponding to our local chunk's parts. + final_returns_chunks = [] + for i, returns_for_seq in enumerate(padded_returns): + returns_for_seq = returns_for_seq[: token_level_rewards[i].size(0)] if cp_size > 1: from vime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + total_len, response_len = total_lengths[i], response_lengths[i] local_returns_chunk = slice_log_prob_with_cp(returns_for_seq, total_len, response_len) else: local_returns_chunk = returns_for_seq @@ -281,7 +448,6 @@ def get_reinforce_plus_plus_returns( def get_reinforce_plus_plus_baseline_advantages( rewards: torch.Tensor, kl: list[torch.Tensor], - loss_masks: list[torch.Tensor], kl_coef: float, ) -> list[torch.Tensor]: """ @@ -293,7 +459,6 @@ def get_reinforce_plus_plus_baseline_advantages( baseline has already been subtracted. kl (list[Tensor]): A list of per-token KL divergence tensors. Used to get the shape for broadcasting. - loss_masks (list[Tensor]): A list of per-token loss masks. kl_coef (float): Coefficient for the KL penalty. Returns: @@ -308,69 +473,6 @@ def get_reinforce_plus_plus_baseline_advantages( return unwhitened_advantages -def get_advantages_and_returns( - total_len: int, - response_len: int, - values: torch.Tensor, - rewards: torch.Tensor, - gamma: float, - lambd: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Function that computes advantages and returns from rewards and values. - Calculated as in the original PPO paper: https://arxiv.org/abs/1707.06347 - Note that rewards may include a KL divergence loss term. - - Advantages looks like this: - Adv1 = R1 + γ * λ * R2 + γ^2 * λ^2 * R3 + ... - - V1 + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ... - - Returns looks like this: - Ret1 = R1 + γ * λ * R2 + γ^2 * λ^2 * R3 + ... - + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ... - - Input: - - values: Tensor of shape (response_size,) - - rewards: Tensor of shape (response_size,) - - Output: - - advantages: Tensor of shape (response_size,) - - returns: Tensor of shape (response_size,) - """ - from megatron.core import mpu - - cp_size = mpu.get_context_parallel_world_size() - if cp_size > 1: - from vime.backends.megatron_utils.cp_utils import all_gather_with_cp - - full_rewards = all_gather_with_cp(rewards, total_len, response_len) - full_values = all_gather_with_cp(values, total_len, response_len) - else: - full_rewards = rewards - full_values = values - - lastgaelam = 0 - advantages_reversed = [] - - for t in reversed(range(response_len)): - nextvalues = full_values[t + 1] if t < response_len - 1 else 0.0 - delta = full_rewards[t] + gamma * nextvalues - full_values[t] - lastgaelam = delta + gamma * lambd * lastgaelam - advantages_reversed.append(lastgaelam) - full_advantages = torch.tensor(advantages_reversed[::-1], dtype=full_values.dtype, device=full_values.device) - full_returns = full_advantages + full_values - - if cp_size > 1: - from vime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp - - advantages = slice_log_prob_with_cp(full_advantages, total_len, response_len) - returns = slice_log_prob_with_cp(full_returns, total_len, response_len) - else: - advantages = full_advantages - returns = full_returns - - return advantages.detach(), returns - - def get_advantages_and_returns_batch( total_lengths, response_lengths, @@ -503,85 +605,45 @@ def vanilla_gae( return full_advantages, full_returns -def chunked_gae( +def chunked_discounted_returns( rewards: torch.Tensor, - values: torch.Tensor, - gamma: float, - lambd: float, + discount: float, chunk_size: int = 128, -): +) -> torch.Tensor: """ - Compute Generalized Advantage Estimation (GAE) using a FlashLinearAttention- - inspired algorithm: parallel prefix scan within chunks and recurrent state - propagation across chunks. + Compute discounted returns using a parallel scan within fixed-size chunks. This reduces the sequential dependency length from O(T) to O(T / chunk_size), while keeping chunk computations fully parallelizable (O(C^2) per chunk). Args: rewards (Tensor): [B, T] reward sequence. - values (Tensor): [B, T] value predictions. The next-value of the final - step is assumed to be zero (standard PPO convention). - gamma (float): discount factor. - lam (float): GAE lambda. + discount (float): Discount factor applied at each step. chunk_size (int): sequence chunk length for parallel scan. Returns: - advantages (Tensor): [B, T] computed advantages. - returns (Tensor): [B, T] advantages + values. + Tensor: [B, T] discounted returns. """ - - # ------------------------------------------------------------------------- - # Validate inputs - # ------------------------------------------------------------------------- - assert rewards.ndim == 2 and values.ndim == 2 + assert rewards.ndim == 2 B, T = rewards.shape - assert values.shape == (B, T) device = rewards.device dtype = rewards.dtype - # ------------------------------------------------------------------------- - # Build δ_t = r_t + γ * V_{t+1} - V_t with V_{T} = 0 - # ------------------------------------------------------------------------- - next_values = torch.cat( - [values[:, 1:], torch.zeros(B, 1, device=device, dtype=dtype)], - dim=1, - ) - deltas = rewards + gamma * next_values - values + # Reformulate the backward recurrence as a forward scan on the reversed + # sequence: S[i] = rewards[i] + discount * S[i - 1]. + rewards_rev = torch.flip(rewards, dims=[1]) - # Reformulate backward GAE as a forward scan on the reversed sequence: - # S[i] = Δ[i] + w * S[i - 1], w = γλ - w = gamma * lambd - deltas_rev = torch.flip(deltas, dims=[1]) # [B, T] - - # ------------------------------------------------------------------------- - # Pad to a multiple of chunk_size - # ------------------------------------------------------------------------- if T % chunk_size != 0: pad = chunk_size - (T % chunk_size) - deltas_rev = F.pad(deltas_rev, (0, pad)) + rewards_rev = F.pad(rewards_rev, (0, pad)) else: pad = 0 - B, T_pad = deltas_rev.shape + B, T_pad = rewards_rev.shape n_chunks = T_pad // chunk_size + rewards_chunks = rewards_rev.view(B, n_chunks, chunk_size) - deltas_chunks = deltas_rev.view(B, n_chunks, chunk_size) - - # ------------------------------------------------------------------------- - # Construct the intra-chunk parallel scan kernel M - # - # For a chunk Δ[0..C-1], we want: - # S_local[t] = sum_{k=0..t} w^(t-k) * Δ[k] - # - # This is implemented as: - # S_local = Δ @ M - # - # where: - # M[i, j] = w^(j - i) if j >= i - # 0 otherwise - # ------------------------------------------------------------------------- idx = torch.arange(chunk_size, device=device) row = idx[:, None] col = idx[None, :] @@ -590,39 +652,25 @@ def chunked_gae( M = torch.zeros(chunk_size, chunk_size, device=device, dtype=dtype) mask = diff >= 0 - if w == 0.0: + if discount == 0.0: M[mask & (diff == 0)] = 1.0 else: - M[mask] = w ** diff[mask].to(dtype) + M[mask] = discount ** diff[mask].to(dtype) - # pow_vec[t] = w^(t+1), used to inject the recurrent state s_prev - if w == 0.0: + if discount == 0.0: pow_vec = torch.zeros(chunk_size, device=device, dtype=dtype) else: - pow_vec = w ** torch.arange(1, chunk_size + 1, device=device, dtype=dtype) + pow_vec = discount ** torch.arange(1, chunk_size + 1, device=device, dtype=dtype) - # ------------------------------------------------------------------------- - # Parallel compute local chunk results (assuming initial state = 0) - # ------------------------------------------------------------------------- - deltas_flat = deltas_chunks.reshape(B * n_chunks, chunk_size) - S_local_flat = deltas_flat @ M + rewards_flat = rewards_chunks.reshape(B * n_chunks, chunk_size) + S_local_flat = rewards_flat @ M S_local_chunks = S_local_flat.view(B, n_chunks, chunk_size) - # Effective length of each chunk (the last chunk may be padded) lengths = [chunk_size] * n_chunks if pad > 0: lengths[-1] = chunk_size - pad - # ------------------------------------------------------------------------- - # Recurrent propagation between chunks - # - # Each chunk contributes: - # S_global[t] = S_local[t] + w^(t+1) * s_prev - # - # And updates: - # s_prev = S_global[last_t] - # ------------------------------------------------------------------------- - S_rev = deltas_rev.new_zeros(B, T_pad) + S_rev = rewards_rev.new_zeros(B, T_pad) s_prev = torch.zeros(B, device=device, dtype=dtype) for c in range(n_chunks): @@ -634,19 +682,46 @@ def chunked_gae( S_global = S_local + s_prev.unsqueeze(1) * pow_vec[:Lc] S_rev[:, start:end] = S_global - s_prev = S_global[:, -1] # state for next chunk + s_prev = S_global[:, -1] - # Remove padding and flip back to original time order if pad > 0: S_rev = S_rev[:, :T] - advantages = torch.flip(S_rev, dims=[1]) + return torch.flip(S_rev, dims=[1]) + + +def chunked_gae( + rewards: torch.Tensor, + values: torch.Tensor, + gamma: float, + lambd: float, + chunk_size: int = 128, +): + """Compute Generalized Advantage Estimation using a chunked scan.""" + assert rewards.ndim == 2 and values.ndim == 2 + B, T = rewards.shape + assert values.shape == (B, T) + + next_values = torch.cat( + [values[:, 1:], torch.zeros(B, 1, device=values.device, dtype=values.dtype)], + dim=1, + ) + deltas = rewards + gamma * next_values - values + advantages = chunked_discounted_returns(deltas, gamma * lambd, chunk_size) returns = advantages + values return advantages, returns -def calculate_log_probs_and_entropy(logits, tokens, tp_group, with_entropy: bool = False, chunk_size: int = -1): +def calculate_log_probs_and_entropy( + logits, + tokens, + tp_group, + with_entropy: bool = False, + chunk_size: int = -1, + log_prob_keep_mask=None, + with_entropy_grad: bool = True, +): logits = logits.contiguous() entropy = None if logits.size(0) != 0: @@ -654,25 +729,36 @@ def calculate_log_probs_and_entropy(logits, tokens, tp_group, with_entropy: bool num_chunks = (logits.size(0) - 1) // chunk_size + 1 logits_chunks = logits.chunk(num_chunks, dim=0) tokens_chunks = tokens.chunk(num_chunks, dim=0) - - if with_entropy: - entropys = [] - for logits_chunk in logits_chunks: - entropy_input = logits_chunk.clone() - entropys.append(compute_entropy_from_logits(entropy_input, tp_group)) - entropy = torch.cat(entropys, dim=0) + mask_chunks = ( + log_prob_keep_mask.chunk(num_chunks, dim=0) if log_prob_keep_mask is not None else [None] * num_chunks + ) log_probs = [] - for tokens_chunk, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True): - log_prob = compute_log_probs(logits_chunk.clone(), tokens_chunk, tp_group) + entropy_chunks = [] + for tokens_chunk, logits_chunk, mask_chunk in zip(tokens_chunks, logits_chunks, mask_chunks, strict=True): + log_prob, entropy_chunk = _calculate_log_probs_and_entropy_chunk( + logits_chunk, + tokens_chunk, + tp_group, + with_entropy=with_entropy, + with_entropy_grad=with_entropy_grad, + log_prob_keep_mask=mask_chunk, + ) log_probs.append(log_prob) + if entropy_chunk is not None: + entropy_chunks.append(entropy_chunk) log_prob = torch.cat(log_probs, dim=0) + if entropy_chunks: + entropy = torch.cat(entropy_chunks, dim=0) else: - if with_entropy: - entropy_input = logits.clone() - entropy = compute_entropy_from_logits(entropy_input, tp_group) - - log_prob = compute_log_probs(logits.clone(), tokens, tp_group) + log_prob, entropy = _calculate_log_probs_and_entropy_chunk( + logits, + tokens, + tp_group, + with_entropy=with_entropy, + with_entropy_grad=with_entropy_grad, + log_prob_keep_mask=log_prob_keep_mask, + ) else: log_prob = logits.new_zeros((0,)) if with_entropy: diff --git a/vime/utils/processing_utils.py b/vime/utils/processing_utils.py index fb652e73d..93f2c9019 100644 --- a/vime/utils/processing_utils.py +++ b/vime/utils/processing_utils.py @@ -16,7 +16,12 @@ def load_tokenizer(name_or_path: str, **kwargs): - return AutoTokenizer.from_pretrained(name_or_path, **kwargs) + tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs) + template_path = Path(name_or_path) / "chat_template.json" + if getattr(tokenizer, "chat_template", None) is None and template_path.is_file(): + with template_path.open() as template_file: + tokenizer.chat_template = json.load(template_file).get("chat_template") + return tokenizer def build_processor_kwargs(multimodal_inputs: dict | None = None) -> dict: @@ -37,6 +42,10 @@ def build_processor_kwargs(multimodal_inputs: dict | None = None) -> dict: else: result[key] = modality_forced.copy() + audio_value = result.get("audio") + if isinstance(audio_value, list): + result["audio"] = [item[0] if isinstance(item, tuple) else item for item in audio_value] + return result @@ -130,12 +139,52 @@ def _extract_images_from_messages(messages): return images +def _load_audio(source): + import soundfile as sf + + audio, sample_rate = sf.read(source, dtype="float32") + if audio.ndim == 2: + audio = audio.mean(axis=1) + if sample_rate != 16000: + from math import gcd + + from scipy.signal import resample_poly + + divisor = gcd(int(sample_rate), 16000) + audio = resample_poly(audio, 16000 // divisor, int(sample_rate) // divisor).astype("float32") + sample_rate = 16000 + return audio, sample_rate + + +def _extract_audios_from_messages(messages): + audios = [] + for message in messages: + content = message.get("content", []) + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict) or item.get("type") != "audio": + continue + audio = item.get("audio") + if audio is None: + audio = item.get("audio_url") + if isinstance(audio, str) and audio.startswith("data:"): + audio = io.BytesIO(base64.b64decode(audio.split(",", 1)[1])) + if isinstance(audio, str) and audio.startswith(("http://", "https://", "file://")): + audios.append(audio) + elif isinstance(audio, str) or hasattr(audio, "read"): + audios.append(_load_audio(audio)) + elif isinstance(audio, tuple): + audios.append(audio) + elif hasattr(audio, "shape"): + audios.append((audio, 16000)) + return audios + + def process_vision_info(prompt, processor): - """Extract PIL images (and videos) from the message list for training. + """Extract image, video, and audio inputs from chat messages.""" + audios = _extract_audios_from_messages(prompt) or None - Tries qwen_vl_utils first (Qwen VL family), falls back to generic - extraction for other models (e.g. GLM-4.6V). - """ try: from qwen_vl_utils import process_vision_info as qwen_process_vision_info @@ -149,7 +198,7 @@ def process_vision_info(prompt, processor): images = _extract_images_from_messages(prompt) or None videos = None - return {"images": images, "videos": videos} + return {"images": images, "videos": videos, "audio": audios} def encode_image_for_rollout_engine(image) -> str: @@ -160,3 +209,56 @@ def encode_image_for_rollout_engine(image) -> str: image.save(buffer, format="PNG") image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/png;base64,{image_base64}" + + +def encode_audio_for_rollout_engine(audio) -> str: + if isinstance(audio, str): + return audio + if not isinstance(audio, tuple) or len(audio) != 2: + raise ValueError(f"Unsupported audio type: {type(audio)}; expected tuple or URL str") + import soundfile as sf + + buffer = io.BytesIO() + sf.write(buffer, *audio, format="WAV", subtype="FLOAT") + return f"data:audio/wav;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}" + + +def encode_video_for_rollout_engine(video) -> str: + if isinstance(video, str): + return video + + import numpy as np + import torch + + if isinstance(video, torch.Tensor): + frames = video.detach().cpu().float() + if frames.dim() == 4 and frames.shape[1] in (1, 3): + frames = frames.permute(0, 2, 3, 1) # (N, H, W, C) + frames = (frames.clamp(0, 1).numpy() * 255).astype(np.uint8) + video = [Image.fromarray(frame) for frame in frames] + + if isinstance(video, list) and video: + encoded = [] + for frame in video: + if not isinstance(frame, Image.Image): + raise ValueError(f"Unsupported video frame type: {type(frame)}") + buffer = io.BytesIO() + frame.save(buffer, format="JPEG") + encoded.append(base64.b64encode(buffer.getvalue()).decode("utf-8")) + return f"data:video/jpeg;base64,{','.join(encoded)}" + + raise ValueError(f"Unsupported video type: {type(video)}") + + +def build_multimodal_messages(prompt: str, multimodal_inputs: dict | None): + multimodal_inputs = multimodal_inputs or {} + content = [{"type": "text", "text": prompt}] + encoders = { + "images": ("image_url", encode_image_for_rollout_engine), + "audio": ("audio_url", encode_audio_for_rollout_engine), + "videos": ("video_url", encode_video_for_rollout_engine), + } + for key, (media_type, encoder) in encoders.items(): + for value in multimodal_inputs.get(key) or []: + content.append({"type": media_type, media_type: {"url": encoder(value)}}) + return [{"role": "user", "content": content}] if len(content) > 1 else None diff --git a/vime/utils/reloadable_process_group.py b/vime/utils/reloadable_process_group.py index fc68b36d4..a09999699 100644 --- a/vime/utils/reloadable_process_group.py +++ b/vime/utils/reloadable_process_group.py @@ -1,16 +1,130 @@ import logging import os from contextlib import contextmanager +from dataclasses import dataclass +from datetime import timedelta +from typing import Any import torch import torch.distributed as dist +from torch.distributed.distributed_c10d import PrefixStore, _get_default_group, _get_default_store -from vime.utils.common import is_npu +from vime.utils import accelerator +from vime.utils.distributed_utils import get_gloo_group, init_gloo_group, set_gloo_group from vime.utils.memory_utils import available_memory, clear_memory, print_memory logger = logging.getLogger(__name__) old_new_group_dict = {} +default_process_group_states = {} + + +@dataclass +class _DefaultProcessGroupState: + backend: str + timeout: timedelta + store: Any + rank: int + world_size: int + generation: int = 0 + accelerator_world_destroyed: bool = False + + +def register_default_process_group(timeout: timedelta) -> None: + """Register the accelerator WORLD group so it can be destroyed and rebuilt. + + Keeping a reference to the rendezvous store is intentional. It keeps the + rank-0 TCPStore alive after ``destroy_process_group()`` and lets every + generation use a fresh PrefixStore namespace, avoiding stale rendezvous + keys when WORLD is recreated repeatedly. + """ + if not dist.is_initialized(): + raise RuntimeError("Cannot register WORLD before torch.distributed is initialized") + + pid = os.getpid() + backend = str(dist.get_backend()) + state = _DefaultProcessGroupState( + backend=backend, + timeout=timeout, + store=_get_default_store(), + rank=dist.get_rank(), + world_size=dist.get_world_size(), + ) + default_process_group_states[pid] = state + logger.info( + "Registered default WORLD process group for reload: backend=%s, rank=%s, world_size=%s", + backend, + state.rank, + state.world_size, + ) + + +def _uses_accelerator_backend(backend: str) -> bool: + return accelerator.is_accelerator_backend(backend) + + +def _new_default_process_group(state: _DefaultProcessGroupState, backend: str) -> None: + state.generation += 1 + store = PrefixStore(f"vime-reloadable-world-{state.generation}-{backend}", state.store) + dist.init_process_group( + backend=backend, + store=store, + rank=state.rank, + world_size=state.world_size, + timeout=state.timeout, + ) + + +def _destroy_default_accelerator_process_group() -> None: + state = default_process_group_states.get(os.getpid()) + if state is None or state.accelerator_world_destroyed or not _uses_accelerator_backend(state.backend): + return + + # Pure PP=4 exposed a teardown ordering deadlock here. Pipeline ranks own + # different overlapping subsets of singleton, embedding, and PP groups, so + # destroying the local wrapper list one group at a time let rank 0 enter + # subgroup reload while another rank was still shutting down. The first + # rank then blocked forever in new_group(), waiting for the others. + # + # Destroying WORLD once makes PyTorch shut down every registered NCCL and + # Gloo backend in its global process-group order. This still releases all + # communicator memory; invalidating the wrappers below only drops stale + # Python handles after their native backends have already been shut down. + dist.barrier(group=get_gloo_group()) + dist.destroy_process_group() + ReloadableProcessGroup.invalidate_process_groups() + set_gloo_group(None) + + _new_default_process_group(state, backend="gloo") + set_gloo_group(_get_default_group()) + state.accelerator_world_destroyed = True + logger.info( + "Destroyed default %s WORLD process group and initialized a temporary Gloo WORLD (generation %s)", + state.backend, + state.generation, + ) + + +def _reload_default_process_group() -> None: + state = default_process_group_states.get(os.getpid()) + if state is None or not state.accelerator_world_destroyed: + return + + # WORLD uses Gloo while the accelerator WORLD is destroyed, so this barrier + # does not recreate an accelerator communicator before all ranks are ready. + dist.barrier() + dist.destroy_process_group() + set_gloo_group(None) + + _new_default_process_group(state, backend=state.backend) + init_gloo_group() + state.accelerator_world_destroyed = False + logger.info( + "Reloaded default WORLD process group with backend %s (generation %s)", + state.backend, + state.generation, + ) + _COMM_MEMORY_CHECK_SKIP_OPS = { "all_gather_into_tensor", @@ -41,9 +155,23 @@ def monkey_patch_torch_dist(): dist.old_new_group = old_new_group def new_group(*args, **kwargs): + explicit_backend = args[2] if len(args) >= 3 else kwargs.get("backend") + backend = str(explicit_backend) if explicit_backend is not None else str(dist.get_backend()) + normalized_backend = accelerator.process_group_backend(backend) if backend == "nccl" else backend + if normalized_backend != backend: + if len(args) >= 3: + args = (*args[:2], normalized_backend, *args[3:]) + else: + kwargs = {**kwargs, "backend": normalized_backend} + backend = normalized_backend + group = old_new_group(*args, **kwargs) - # skip none nccl group. - if len(args) >= 3 and args[2] == "gloo" or "backend" in kwargs and kwargs["backend"] == "gloo": + + # Before WORLD is registered, preserve the historical behavior of + # leaving CPU groups and singleton groups untouched. Afterwards every + # cached subgroup must be reloadable because destroying WORLD + # invalidates all of them, including Gloo and singleton groups. + if backend == "gloo" and pid not in default_process_group_states: return group # Get ranks from arguments @@ -55,10 +183,22 @@ def new_group(*args, **kwargs): # If no ranks specified, use all ranks in world ranks = list(range(dist.get_world_size())) - if len(ranks) == 1: + # Historically singleton groups were left unwrapped because they do + # not own a useful communicator. Once WORLD itself is destroyed, + # however, PyTorch invalidates *every* registered subgroup, including + # singleton groups cached by Megatron. Wrap them for actors that have + # registered a reloadable WORLD so those cached references remain + # usable after wake-up. Preserve the old behavior for other callers. + if len(ranks) == 1 and pid not in default_process_group_states: return group - group = ReloadableProcessGroup(group, ranks) + group = ReloadableProcessGroup( + group, + ranks, + creation_args=args, + creation_kwargs=kwargs, + backend=backend, + ) return group dist.new_group = new_group @@ -104,10 +244,14 @@ def new_function(*args, **kwargs): dist.reduce_scatter = get_new_comm_function(dist.reduce_scatter) dist.reduce_scatter_tensor = get_new_comm_function(dist.reduce_scatter_tensor, "reduce_scatter_tensor") dist.scatter = get_new_comm_function(dist.scatter) + dist.scatter_object_list = get_new_comm_function(dist.scatter_object_list) dist.gather = get_new_comm_function(dist.gather) + dist.gather_object = get_new_comm_function(dist.gather_object) dist.barrier = get_new_comm_function(dist.barrier, "barrier") dist.send = get_new_comm_function(dist.send) + dist.send_object_list = get_new_comm_function(dist.send_object_list) dist.recv = get_new_comm_function(dist.recv) + dist.recv_object_list = get_new_comm_function(dist.recv_object_list) dist._coalescing_manager = get_new_comm_function(dist._coalescing_manager) # p2p @@ -141,7 +285,7 @@ def convert(arg): class ReloadableProcessGroup(torch.distributed.ProcessGroup): GROUPS = {} - def __init__(self, group, ranks): + def __init__(self, group, ranks, *, creation_args=(), creation_kwargs=None, backend="nccl"): super().__init__( rank=dist.get_rank(group), size=dist.get_world_size(group), @@ -149,6 +293,9 @@ def __init__(self, group, ranks): self.group = group self.group_info = { "ranks": ranks, + "args": tuple(creation_args), + "kwargs": dict(creation_kwargs or {}), + "backend": backend, } pid = os.getpid() if pid not in ReloadableProcessGroup.GROUPS: @@ -175,19 +322,35 @@ def destroy_process_groups(): del reloadable_group.group reloadable_group.group = None + @staticmethod + def invalidate_process_groups(): + """Drop handles after destroying WORLD, which already shut down every subgroup.""" + pid = os.getpid() + for reloadable_group in ReloadableProcessGroup.GROUPS.get(pid, []): + reloadable_group.group = None + @staticmethod def reload_process_groups(): pid = os.getpid() reloadable_groups = ReloadableProcessGroup.GROUPS.get(pid, []) - logger.info(f"Reloading {len(reloadable_groups)} process groups in pid {pid}") + backend_counts = {} + for reloadable_group in reloadable_groups: + backend = reloadable_group.group_info["backend"] + backend_counts[backend] = backend_counts.get(backend, 0) + 1 + logger.info( + "Reloading %s process groups in pid %s: %s", + len(reloadable_groups), + pid, + backend_counts, + ) old_new_group = old_new_group_dict.get(pid) - backend = "nccl" - if is_npu(): - backend = "hccl" for reloadable_group in reloadable_groups: if reloadable_group.group is not None: continue - group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend=backend) + group = old_new_group( + *reloadable_group.group_info["args"], + **reloadable_group.group_info["kwargs"], + ) reloadable_group.group = group def rank(self) -> int: @@ -224,6 +387,9 @@ def _fwd_query(self, method, *args, **kwargs): def barrier(self, *a, **kw): return self._fwd("barrier", *a, **kw) + def monitored_barrier(self, *a, **kw): + return self._fwd("monitored_barrier", *a, **kw) + def broadcast(self, *a, **kw): return self._fwd("broadcast", *a, **kw) @@ -248,6 +414,12 @@ def allgather_coalesced(self, *a, **kw): def allgather_into_tensor_coalesced(self, *a, **kw): return self._fwd("allgather_into_tensor_coalesced", *a, **kw) + def all_gather_single(self, *a, **kw): + return self._fwd("all_gather_single", *a, **kw) + + def all_gather_single_coalesced(self, *a, **kw): + return self._fwd("all_gather_single_coalesced", *a, **kw) + def gather(self, *a, **kw): return self._fwd("gather", *a, **kw) @@ -257,6 +429,12 @@ def scatter(self, *a, **kw): def reduce_scatter(self, *a, **kw): return self._fwd("reduce_scatter", *a, **kw) + def reduce_scatter_single(self, *a, **kw): + return self._fwd("reduce_scatter_single", *a, **kw) + + def reduce_scatter_single_coalesced(self, *a, **kw): + return self._fwd("reduce_scatter_single_coalesced", *a, **kw) + def _reduce_scatter_base(self, *a, **kw): return self._fwd("_reduce_scatter_base", *a, **kw) @@ -303,12 +481,17 @@ def bound_device_id(self, dev): def destroy_process_groups(): - """Destroy all reloadable process groups.""" - ReloadableProcessGroup.destroy_process_groups() + """Destroy registered subgroups and replace accelerator WORLD with a temporary Gloo WORLD.""" + state = default_process_group_states.get(os.getpid()) + if state is not None and not state.accelerator_world_destroyed and _uses_accelerator_backend(state.backend): + _destroy_default_accelerator_process_group() + else: + ReloadableProcessGroup.destroy_process_groups() def reload_process_groups(): - """Reload all reloadable process groups.""" + """Restore accelerator WORLD and recreate all registered subgroups.""" + _reload_default_process_group() ReloadableProcessGroup.reload_process_groups() diff --git a/vime/utils/rocm_checkpoint_writer.py b/vime/utils/rocm_checkpoint_writer.py new file mode 100644 index 000000000..7a8a1be2c --- /dev/null +++ b/vime/utils/rocm_checkpoint_writer.py @@ -0,0 +1,27 @@ +import torch +from megatron.core.dist_checkpointing.strategies.filesystem_async import FileSystemWriterAsync + + +class ROCmFileSystemWriterAsync(FileSystemWriterAsync): + """ + FileSystemWriterAsync wrapper for ROCm compatibility. + + On ROCm/HIP, using non_blocking=True causes tensors to be stored in pinned memory, + which triggers segmentation faults when forking subprocesses afterward. + """ + + @staticmethod + def preload_tensors(*args, **kwargs): + # Change argument non_blocking to False on HIP platform + # The tensors will be stored in pinned memory if non_blocking=True + # Currently on the ROCm platform, forking a subprocess afterward + # with pinned_memory=True will trigger segmentation fault + if torch.version.hip: + print("HIP/ROCm detected: setting non_blocking=False in preload_tensors") + if "non_blocking" in kwargs: + kwargs["non_blocking"] = False + elif len(args) > 1 and isinstance(args[-1], bool): + # non_blocking is typically the last argument + args = args[:-1] + (False,) + + return FileSystemWriterAsync.preload_tensors(*args, **kwargs) diff --git a/vime/utils/routing_replay.py b/vime/utils/routing_replay.py index 864728166..96c199dca 100644 --- a/vime/utils/routing_replay.py +++ b/vime/utils/routing_replay.py @@ -1,8 +1,11 @@ import os import torch +from vime.utils import accelerator + ROUTING_REPLAY = None +ORDERED_TOPK_CAPTURE_ROUTER = None def set_routing_replay(replay): @@ -10,8 +13,74 @@ def set_routing_replay(replay): ROUTING_REPLAY = replay +def _set_ordered_topk_capture_router(router): + global ORDERED_TOPK_CAPTURE_ROUTER + ORDERED_TOPK_CAPTURE_ROUTER = router + + +def _capture_ordered_topk(top_indices): + """Keep the current router's exact top-k order until its MoE combine.""" + router = ORDERED_TOPK_CAPTURE_ROUTER + if router is not None: + router._vime_ordered_topk_indices = top_indices + + +def consume_ordered_topk(module): + """Return and release the current forward's ordered top-k indices.""" + return module.__dict__.pop("_vime_ordered_topk_indices", None) + + +def register_ordered_topk_capture(module): + """Capture one forward's VLLM-compatible top-k order without R3.""" + if getattr(module, "_vime_ordered_topk_capture_registered", False): + return + + def pre_forward_hook(patched_module, *args, **kwargs): + del args, kwargs + _set_ordered_topk_capture_router(patched_module) + + def forward_hook(patched_module, *args, **kwargs): + del args, kwargs + if ORDERED_TOPK_CAPTURE_ROUTER is patched_module: + _set_ordered_topk_capture_router(None) + + module.register_forward_pre_hook(pre_forward_hook) + module.register_forward_hook(forward_hook) + module._vime_ordered_topk_capture_registered = True + + +def _compute_topk_for_current_router( + old_compute_topk, + scores, + topk, + num_groups=None, + group_topk=None, +): + # VLLM's deterministic DeepSeek/GLM biased top-k uses + # torch.topk(..., sorted=False). Megatron's local compute_topk uses the + # default sorted=True. The selected expert set is the same, but the + # low-latency-compatible owner reduction consumes experts in top-k column + # order, so the difference changes BF16 accumulation before any route + # actually diverges. + # + # Only override routers registered by the DeepEP alignment bridge, and only for the + # non-grouped Megatron path used by GLM-5 (n_group=topk_group=1 in VLLM, + # represented as no group limit in Megatron). Other training paths retain + # Megatron's original semantics. + if ORDERED_TOPK_CAPTURE_ROUTER is not None and not group_topk: + return torch.topk(scores, k=topk, dim=1, sorted=False) + + return old_compute_topk( + scores, + topk, + num_groups=num_groups, + group_topk=group_topk, + ) + + class RoutingReplay: all_routing_replays = [] + lazy_resources = [] def __init__(self): self.forward_index = 0 @@ -20,20 +89,53 @@ def __init__(self): RoutingReplay.all_routing_replays.append(self) def record(self, top_indices): - # offload top_indices to CPU pinned memory + if hasattr(top_indices, "materialize_for_routing_replay"): + self.top_indices_list.append(top_indices) + return + if top_indices.device.type == "cpu" and top_indices.is_pinned() and top_indices.is_contiguous(): + self.top_indices_list.append(top_indices) + return + + # Compact non-contiguous layer views so they do not retain the full + # all-layer routed-experts tensor once per pipeline stage. buf = torch.empty_like(top_indices, device="cpu", pin_memory=True) buf.copy_(top_indices) self.top_indices_list.append(buf) + @staticmethod + def assert_all_consumed() -> None: + errors = [] + for replay_idx, replay in enumerate(RoutingReplay.all_routing_replays): + expected = len(replay.top_indices_list) + if replay.forward_index != expected or replay.backward_index != expected: + errors.append( + f"router={replay_idx}: forward={replay.forward_index}, " + f"backward={replay.backward_index}, recorded={expected}" + ) + if errors: + raise RuntimeError("R3 routing replay was not consumed exactly once: " + "; ".join(errors[:16])) + def pop_forward(self): top_indices = self.top_indices_list[self.forward_index] self.forward_index += 1 - return top_indices.to(torch.cuda.current_device()) + if hasattr(top_indices, "materialize_for_routing_replay"): + return top_indices.materialize_for_routing_replay("forward") + return top_indices.to( + accelerator.current_device(), + dtype=torch.int32, + non_blocking=top_indices.is_pinned(), + ) def pop_backward(self): top_indices = self.top_indices_list[self.backward_index] self.backward_index += 1 - return top_indices.to(torch.cuda.current_device()) + if hasattr(top_indices, "materialize_for_routing_replay"): + return top_indices.materialize_for_routing_replay("backward") + return top_indices.to( + accelerator.current_device(), + dtype=torch.int32, + non_blocking=top_indices.is_pinned(), + ) def clear(self): self.forward_index = 0 @@ -47,21 +149,45 @@ def clear_forward(self): def clear_all(): for replay in RoutingReplay.all_routing_replays: replay.clear() + for resource in RoutingReplay.lazy_resources: + resource.close() + RoutingReplay.lazy_resources = [] @staticmethod def clear_all_forward(): for replay in RoutingReplay.all_routing_replays: replay.clear_forward() + @staticmethod + def register_lazy_resource(resource): + RoutingReplay.lazy_resources.append(resource) + + @staticmethod + def begin_lazy_pass(release_stage): + for resource in RoutingReplay.lazy_resources: + resource.begin_pass(release_stage) + def get_routing_replay_compute_topk(old_compute_topk): def compute_topk(scores, topk, num_groups=None, group_topk=None): if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": routing_replay_stage = os.environ["ROUTING_REPLAY_STAGE"] if routing_replay_stage == "fallthrough": - return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) - if routing_replay_stage == "record": - probs, top_indices = old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) + probs, top_indices = _compute_topk_for_current_router( + old_compute_topk, + scores, + topk, + num_groups=num_groups, + group_topk=group_topk, + ) + elif routing_replay_stage == "record": + probs, top_indices = _compute_topk_for_current_router( + old_compute_topk, + scores, + topk, + num_groups=num_groups, + group_topk=group_topk, + ) ROUTING_REPLAY.record(top_indices) elif routing_replay_stage == "replay_forward": top_indices = ROUTING_REPLAY.pop_forward() @@ -75,9 +201,17 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): top_indices.shape[0] == scores.shape[0] and top_indices.shape[1] == topk ), f"top_indices shape {top_indices.shape} does not match scores shape {scores.shape} and topk {topk}" probs = scores.gather(1, top_indices) - return probs, top_indices else: - return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) + probs, top_indices = _compute_topk_for_current_router( + old_compute_topk, + scores, + topk, + num_groups=num_groups, + group_topk=group_topk, + ) + + _capture_ordered_topk(top_indices) + return probs, top_indices return compute_topk diff --git a/vime/utils/seqlen_balancing.py b/vime/utils/seqlen_balancing.py index 5736d8850..57cf47f3c 100644 --- a/vime/utils/seqlen_balancing.py +++ b/vime/utils/seqlen_balancing.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import heapq @@ -123,26 +122,6 @@ def __repr__(self) -> str: return partitions -def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool): - bias = sum(seqlen_list) + 1 if equal_size else 0 - sorted_seqlen = [(seqlen + bias, i) for i, seqlen in enumerate(seqlen_list)] - partitions = [[] for _ in range(k_partitions)] - partition_sums = [0 for _ in range(k_partitions)] - for seqlen, i in sorted_seqlen: - min_idx = None - for j in range(k_partitions): - if min_idx is None or partition_sums[j] < partition_sums[min_idx]: - min_idx = j - partitions[min_idx].append(i) - partition_sums[min_idx] += seqlen - if equal_size: - for _i, partition in enumerate(partitions): - assert len(partition) * k_partitions == len( - seqlen_list - ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" - return partitions - - def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, equal_size: bool): """get order of seq lengths to make partitions balanced, this is used in balacing sum of seqlength across dp ranks and microbatches @@ -227,12 +206,3 @@ def expand_bins_by_splitting(bins: list[list[int]], target_count: int, lengths) left, right = _split_bin_by_tokens(bins[idx], lengths) bins[idx] = left bins.append(right) - - -def get_reverse_idx(idx_map): - reverse_idx_map = copy.deepcopy(idx_map) - - for i, idx in enumerate(idx_map): - reverse_idx_map[idx] = i - - return reverse_idx_map diff --git a/vime/utils/tensor_backper.py b/vime/utils/tensor_backper.py index 2fc2a6359..023971de9 100644 --- a/vime/utils/tensor_backper.py +++ b/vime/utils/tensor_backper.py @@ -1,47 +1,16 @@ -from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Callable, Iterable import torch -_SourceGetter = Callable[[], Iterable[tuple[str, torch.Tensor]]] +from vime.utils import accelerator +_SourceGetter = Callable[[], Iterable[tuple[str, torch.Tensor]]] -class TensorBackuper(ABC): - @staticmethod - def create(source_getter, single_tag): - if single_tag is None: - return _TensorBackuperNormal(source_getter=source_getter) - else: - return _TensorBackuperNoop(source_getter=source_getter, single_tag=single_tag) +class TensorBackuper: def __init__(self, source_getter: _SourceGetter): self._source_getter = source_getter - - @property - @abstractmethod - def backup_tags(self): - raise NotImplementedError - - @abstractmethod - def get(self, tag: str): - raise NotImplementedError - - @abstractmethod - def backup(self, tag: str): - raise NotImplementedError - - def copy(self, *, src_tag: str, dst_tag: str): - raise NotImplementedError - - @abstractmethod - def restore(self, tag: str): - raise NotImplementedError - - -class _TensorBackuperNormal(TensorBackuper): - def __init__(self, source_getter): - super().__init__(source_getter=source_getter) self._backups: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) @property @@ -58,7 +27,7 @@ def backup(self, tag: str) -> None: if name not in backup_dict: backup_dict[name] = torch.empty_like(param, device=torch.device("cpu"), pin_memory=True) backup_dict[name].copy_(param.detach(), non_blocking=True) - torch.cuda.synchronize() + accelerator.synchronize() @torch.no_grad() def copy(self, *, src_tag: str, dst_tag: str): @@ -71,45 +40,4 @@ def restore(self, tag: str) -> None: for name, param in self._source_getter(): assert name in backup_dict param.copy_(backup_dict[name], non_blocking=True) - torch.cuda.synchronize() - - -class _TensorBackuperNoop(TensorBackuper): - def __init__(self, source_getter, single_tag): - super().__init__(source_getter=source_getter) - self._single_tag = single_tag - # Sanity check for safety - self._backup_hash_dict = None - - @property - def backup_tags(self): - return [self._single_tag] - - def get(self, tag: str): - ans = dict(self._source_getter()) - ans = {k: v.detach() for k, v in ans.items()} - assert _compute_hash_dict(ans) == self._backup_hash_dict - return ans - - def backup(self, tag: str) -> None: - assert tag == self._single_tag - self._backup_hash_dict = _compute_hash_dict(dict(self._source_getter())) - torch.cuda.synchronize() - - def restore(self, tag: str) -> None: - assert tag == self._single_tag - assert _compute_hash_dict(dict(self._source_getter())) == self._backup_hash_dict - torch.cuda.synchronize() - - -def _compute_hash_dict(tensors: dict[str, torch.Tensor]): - return {k: _compute_hash_tensor(v) for k, v in tensors.items()} - - -def _compute_hash_tensor(x: torch.Tensor): - # Not a real/good hash, but pretty fast - x = x.contiguous() - x = x.view(-1) - x = x.view(torch.uint32) - x = x.sum() - return x.item() + accelerator.synchronize() diff --git a/vime/utils/train_dump_utils.py b/vime/utils/train_dump_utils.py deleted file mode 100644 index 4b5cbc6a5..000000000 --- a/vime/utils/train_dump_utils.py +++ /dev/null @@ -1,22 +0,0 @@ -import logging -from pathlib import Path - -import torch - -logger = logging.getLogger(__name__) - - -def save_debug_train_data(args, *, rollout_id, rollout_data): - if (path_template := args.save_debug_train_data) is not None: - rank = torch.distributed.get_rank() - path = Path(path_template.format(rollout_id=rollout_id, rank=rank)) - logger.info(f"Save debug train data to {path}") - path.parent.mkdir(parents=True, exist_ok=True) - torch.save( - dict( - rollout_id=rollout_id, - rank=rank, - rollout_data=rollout_data, - ), - path, - ) diff --git a/vime/utils/train_metric_utils.py b/vime/utils/train_metric_utils.py deleted file mode 100644 index 0782cb2b7..000000000 --- a/vime/utils/train_metric_utils.py +++ /dev/null @@ -1,48 +0,0 @@ -import logging -from argparse import Namespace -from collections.abc import Callable -from copy import deepcopy - -from vime.utils import logging_utils -from vime.utils.metric_utils import compute_rollout_step -from vime.utils.timer import Timer - -logger = logging.getLogger(__name__) - - -def log_perf_data_raw( - rollout_id: int, args: Namespace, is_primary_rank: bool, compute_total_fwd_flops: Callable -) -> None: - timer_instance = Timer() - log_dict_raw = deepcopy(timer_instance.log_dict()) - timer_instance.reset() - - if not is_primary_rank: - return - - log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()} - - if ("perf/actor_train_time" in log_dict) and (compute_total_fwd_flops is not None): - total_fwd_flops = compute_total_fwd_flops(seq_lens=timer_instance.seq_lens) - - if "perf/log_probs_time" in log_dict: - log_dict["perf/log_probs_tflops"] = total_fwd_flops / log_dict["perf/log_probs_time"] - - if "perf/ref_log_probs_time" in log_dict: - log_dict["perf/ref_log_probs_tflops"] = total_fwd_flops / log_dict["perf/ref_log_probs_time"] - - if log_dict["perf/actor_train_time"] > 0: - log_dict["perf/actor_train_tflops"] = 3 * total_fwd_flops / log_dict["perf/actor_train_time"] - log_dict["perf/actor_train_tok_per_s"] = sum(timer_instance.seq_lens) / log_dict["perf/actor_train_time"] - - if "perf/train_wait_time" in log_dict and "perf/train_time" in log_dict: - total_time = log_dict["perf/train_wait_time"] + log_dict["perf/train_time"] - if total_time > 0: - log_dict["perf/step_time"] = total_time - log_dict["perf/wait_time_ratio"] = log_dict["perf/train_wait_time"] / total_time - - logger.info(f"perf {rollout_id}: {log_dict}") - - step = compute_rollout_step(args, rollout_id) - log_dict["rollout/step"] = step - logging_utils.log(args, log_dict, step_key="rollout/step") diff --git a/vime/utils/types.py b/vime/utils/types.py index 7e2a45bf3..45fa07697 100644 --- a/vime/utils/types.py +++ b/vime/utils/types.py @@ -4,6 +4,91 @@ import torch +from vime.utils.misc import decode_int32_meta_array + +_TOP_P_TOKEN_ID_META_KEYS = ("top_p_token_ids", "top_p_kept_token_ids") +_TOP_P_TOKEN_OFFSET_META_KEYS = ("top_p_token_offsets", "top_p_kept_token_offsets") + + +def _extract_rollout_top_p_token_data( + meta_info: dict[str, Any], + *, + expected_num_tokens: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor] | None: + token_ids = decode_int32_meta_array(meta_info, _TOP_P_TOKEN_ID_META_KEYS) + offsets = decode_int32_meta_array(meta_info, _TOP_P_TOKEN_OFFSET_META_KEYS) + if token_ids is None and offsets is None: + return None + if token_ids is None or offsets is None: + raise ValueError("VLLM top-p token replay must include both token ids and offsets.") + if offsets.numel() == 0 or int(offsets[0]) != 0: + raise ValueError(f"VLLM top-p token offsets must start with 0, got {offsets[:1].tolist()}.") + if int(offsets[-1]) != token_ids.numel(): + raise ValueError( + "VLLM top-p token ids/offsets mismatch: " + f"offsets[-1]={int(offsets[-1])}, len(token_ids)={token_ids.numel()}." + ) + if expected_num_tokens is not None and offsets.numel() != expected_num_tokens + 1: + raise ValueError( + "VLLM top-p token offsets length must equal generated token count + 1: " + f"len(offsets)={offsets.numel()}, generated={expected_num_tokens}." + ) + return token_ids, offsets + + +def _merge_rollout_top_p_token_data( + base_token_ids: list[int] | torch.Tensor | None, + base_offsets: list[int] | torch.Tensor | None, + token_ids: torch.Tensor, + offsets: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + base_token_ids = torch.as_tensor([] if base_token_ids is None else base_token_ids, dtype=torch.int32).reshape(-1) + base_offsets = torch.as_tensor([0] if base_offsets is None else base_offsets, dtype=torch.int32).reshape(-1) + base_offset = int(base_offsets[-1]) + return ( + torch.cat([base_token_ids, token_ids]), + torch.cat([base_offsets, offsets[1:] + base_offset]), + ) + + +def _pad_rollout_top_p_offsets( + token_ids: list[int] | torch.Tensor | None, + offsets: list[int] | torch.Tensor | None, + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if offsets is None or token_ids is None: + raise ValueError("Cannot append empty top-p spans without existing token ids and offsets.") + if num_tokens < 0: + raise ValueError(f"num_tokens must be non-negative, got {num_tokens}.") + token_ids = torch.as_tensor(token_ids, dtype=torch.int32).reshape(-1) + offsets = torch.as_tensor(offsets, dtype=torch.int32).reshape(-1) + if offsets.numel() == 0: + raise ValueError("Cannot append empty top-p spans to empty offsets.") + if num_tokens == 0: + return token_ids, offsets + empty_offsets = offsets.new_full((num_tokens,), int(offsets[-1])) + return token_ids, torch.cat([offsets, empty_offsets]) + + +def _to_int_list(tokens) -> list[int]: + if tokens is None: + return [] + if torch.is_tensor(tokens): + return [int(token) for token in tokens.detach().cpu().reshape(-1).tolist()] + return [int(token) for token in tokens] + + +def _to_float_list(values) -> list[float] | None: + if values is None: + return None + if torch.is_tensor(values): + return [float(value) for value in values.detach().cpu().reshape(-1).tolist()] + return [float(value) for value in values] + + +def _numel(value) -> int: + return int(torch.as_tensor(value).reshape(-1).numel()) + @dataclass class Sample: @@ -24,6 +109,8 @@ class Sample: tokens: list[int] = field(default_factory=list) multimodal_inputs: dict[str, Any] | None = None # raw multimodal data, e.g. images, videos, etc. multimodal_train_inputs: dict[str, Any] | None = None # processed multimodal data, e.g. pixel_values, etc. + multimodal_train_input_id: str | None = None + apply_chat_template_kwargs: dict = field(default_factory=dict) # response response: str = "" response_length: int = 0 @@ -32,7 +119,11 @@ class Sample: loss_mask: list[int] | None = None weight_versions: list[str] = field(default_factory=list) rollout_log_probs: list[float] | None = None # Log probabilities from rollout engine - rollout_routed_experts: list[list[int]] | None = None # Routed experts from rollout engine + # Ragged top-p nucleus token ids replayed from rollout sampling. For response + # token i, kept ids are rollout_top_p_token_ids[offsets[i]:offsets[i + 1]]. + rollout_top_p_token_ids: list[int] | torch.Tensor | None = None + rollout_top_p_token_offsets: list[int] | torch.Tensor | None = None + rollout_routed_experts: list[list[int]] | torch.Tensor | None = None # Routed experts from rollout engine remove_sample: bool = False teacher_log_probs: list[float] | None = None # Log probabilities from teacher model for OPD @@ -50,10 +141,11 @@ class Status(Enum): metadata: dict = field(default_factory=dict) generate_function_path: str | None = None + custom_rm_path: str | None = None # metadata used during training, e.g., what loss to use for this sample. train_metadata: dict | None = None - # Session ID for consistent hashing routing (used when router policy is consistent_hashing) + # Session ID for consistent hashing routing (used when router policy is consistent_hash). session_id: str | None = None non_generation_time: float = 0.0 # time spent in non-generation steps @@ -158,12 +250,154 @@ def get_reward_value(self, args) -> float: def effective_response_length(self): return sum(self.loss_mask) if self.loss_mask is not None else self.response_length - def update_from_meta_info(self, args, meta_info: dict): + def append_response_tokens( + self, + args=None, + *, + tokens=None, + log_probs=None, + trainable: bool = True, + meta_info: dict | None = None, + text: str | None = None, + update_terminal_info: bool = True, + ): """ - Update the sample with new information from meta_info returned by the rollout engine. - And extract + Append response-side tokens and keep training metadata aligned. + + Model-generated tokens should pass ``trainable=True`` plus VLLM + ``meta_info`` and log probabilities. Tool/environment tokens should pass + ``trainable=False``; they receive loss-mask zeros and empty top-p spans + when top-p replay is active. """ - if args.vllm_speculative_config: + tokens = _to_int_list(tokens) + log_probs = _to_float_list(log_probs) + if log_probs is not None and len(log_probs) != len(tokens): + raise ValueError(f"log_probs length {len(log_probs)} != tokens length {len(tokens)}") + if tokens and trainable and log_probs is None: + raise ValueError("trainable response tokens require rollout log probabilities.") + if tokens and not trainable: + if log_probs is not None: + raise ValueError("non-trainable response tokens should not pass rollout log probabilities.") + log_probs = [0.0] * len(tokens) + + if text is not None: + self.response += text + + previous_response_length = self.response_length + if tokens: + self.tokens += tokens + self.response_length += len(tokens) + if self.loss_mask is None: + self.loss_mask = [1] * previous_response_length + self.loss_mask += [1 if trainable else 0] * len(tokens) + + if log_probs is not None: + if self.rollout_log_probs is None: + if trainable and previous_response_length: + raise ValueError( + "Cannot append trainable rollout log probabilities to a sample with existing response " + "tokens but no existing rollout_log_probs." + ) + self.rollout_log_probs = [0.0] * previous_response_length + self.rollout_log_probs += log_probs + + should_pad_top_p = bool(tokens and not trainable) + if meta_info is not None or should_pad_top_p: + self._apply_meta_info( + args, + meta_info or {}, + new_token_count=len(tokens), + pad_missing_top_p=should_pad_top_p, + update_terminal_info=update_terminal_info, + ) + + self._validate_response_metadata_lengths() + + def _apply_meta_info( + self, + args, + meta_info: dict, + *, + new_token_count: int = 0, + pad_missing_top_p: bool = False, + update_terminal_info: bool = True, + ) -> None: + applied_top_p_data = False + if new_token_count: + top_p_data = _extract_rollout_top_p_token_data(meta_info, expected_num_tokens=new_token_count) + if top_p_data is not None: + applied_top_p_data = True + base_token_ids, base_offsets = self.rollout_top_p_token_ids, self.rollout_top_p_token_offsets + if base_token_ids is None and base_offsets is None: + self.rollout_top_p_token_ids, self.rollout_top_p_token_offsets = top_p_data + else: + self.rollout_top_p_token_ids, self.rollout_top_p_token_offsets = _merge_rollout_top_p_token_data( + base_token_ids, + base_offsets, + *top_p_data, + ) + + if ( + pad_missing_top_p + and new_token_count + and self.rollout_top_p_token_offsets is not None + and not applied_top_p_data + ): + self.rollout_top_p_token_ids, self.rollout_top_p_token_offsets = _pad_rollout_top_p_offsets( + self.rollout_top_p_token_ids, + self.rollout_top_p_token_offsets, + new_token_count, + ) + + routed_experts = decode_int32_meta_array(meta_info, "routed_experts") + if routed_experts is not None: + if args is None: + raise ValueError("args is required to decode routed experts metadata.") + routed_experts_start_len = int(meta_info.get("routed_experts_start_len", 0) or 0) + if routed_experts_start_len < 0: + raise ValueError( + f"VLLM routed_experts_start_len must be non-negative, got {routed_experts_start_len}." + ) + expected_rows = max(0, len(self.tokens) - 1 - routed_experts_start_len) + expected_numel = expected_rows * args.num_layers * args.moe_router_topk + if routed_experts.numel() != expected_numel: + raise ValueError( + "VLLM routed_experts element count does not match sample tokens: " + f"got={routed_experts.numel()}, expected={expected_numel} " + f"(tokens={len(self.tokens)}, routed_experts_start_len={routed_experts_start_len}, " + f"num_layers={args.num_layers}, " + f"moe_router_topk={args.moe_router_topk})." + ) + routed_experts = routed_experts.reshape( + expected_rows, + args.num_layers, + args.moe_router_topk, + ) + if routed_experts_start_len == 0: + self.rollout_routed_experts = routed_experts + else: + existing = self.rollout_routed_experts + if existing is None: + raise ValueError( + "Cannot append partial routed experts without existing routed experts " + f"(routed_experts_start_len={routed_experts_start_len})." + ) + if not torch.is_tensor(existing): + existing = torch.as_tensor(existing, dtype=routed_experts.dtype) + if existing.shape[0] < routed_experts_start_len: + raise ValueError( + "Existing routed experts shorter than routed_experts_start_len: " + f"existing_rows={existing.shape[0]}, routed_experts_start_len={routed_experts_start_len}." + ) + self.rollout_routed_experts = torch.cat( + [existing[:routed_experts_start_len], routed_experts], + dim=0, + ) + + if not update_terminal_info or "finish_reason" not in meta_info: + return + + if getattr(args, "vllm_speculative_config", None): # cannot directly use spec info from vLLM because of partial rollout. self.spec_info.add(meta_info=meta_info) @@ -181,6 +415,33 @@ def update_from_meta_info(self, args, meta_info: dict): case "stop": self.status = Sample.Status.COMPLETED + def _validate_response_metadata_lengths(self): + if self.loss_mask is not None and len(self.loss_mask) != self.response_length: + raise ValueError(f"loss_mask length {len(self.loss_mask)} != response_length {self.response_length}") + + if self.rollout_log_probs is not None and len(self.rollout_log_probs) != self.response_length: + raise ValueError( + f"rollout_log_probs length {len(self.rollout_log_probs)} != response_length {self.response_length}" + ) + + if self.rollout_top_p_token_ids is None and self.rollout_top_p_token_offsets is None: + return + if self.rollout_top_p_token_ids is None or self.rollout_top_p_token_offsets is None: + raise ValueError("rollout top-p replay must include both token ids and offsets.") + + offsets = torch.as_tensor(self.rollout_top_p_token_offsets, dtype=torch.int32).reshape(-1) + if offsets.numel() != self.response_length + 1: + raise ValueError( + "rollout_top_p_token_offsets length must equal response_length + 1: " + f"len(offsets)={offsets.numel()}, response_length={self.response_length}." + ) + token_id_count = _numel(self.rollout_top_p_token_ids) + if int(offsets[-1]) != token_id_count: + raise ValueError( + "rollout top-p token ids/offsets mismatch: " + f"offsets[-1]={int(offsets[-1])}, len(token_ids)={token_id_count}." + ) + @dataclass(frozen=True) class ParamInfo: diff --git a/vime_plugins/mbridge/__init__.py b/vime_plugins/mbridge/__init__.py deleted file mode 100644 index 9263cbe90..000000000 --- a/vime_plugins/mbridge/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -from .deepseek_v32 import DeepseekV32Bridge -from .glm4 import GLM4Bridge -from .glm4moe import GLM4MoEBridge -from .glm4moe_lite import GLM4MoELiteBridge -from .gpt_oss import GptOssBridge -from .mimo import MimoBridge -from .minimax_m2 import MiniMaxM2Bridge -from .qwen3_5 import Qwen3_5Bridge -from .qwen3_next import Qwen3NextBridge - -__all__ = [ - "GLM4Bridge", - "GLM4MoEBridge", - "GLM4MoELiteBridge", - "GptOssBridge", - "MiniMaxM2Bridge", - "Qwen3NextBridge", - "Qwen3_5Bridge", - "MimoBridge", - "DeepseekV32Bridge", -] diff --git a/vime_plugins/mbridge/deepseek_v32.py b/vime_plugins/mbridge/deepseek_v32.py deleted file mode 100644 index d45fd40fe..000000000 --- a/vime_plugins/mbridge/deepseek_v32.py +++ /dev/null @@ -1,80 +0,0 @@ -import torch - -from mbridge.core import register_model -from mbridge.models import DeepseekV3Bridge - - -@register_model("deepseek_v32") -class DeepseekV32Bridge(DeepseekV3Bridge): - - def __init__(self, hf_config, **kwargs): - # transformers 5.x stores rope_theta inside rope_parameters dict, - # but DeepseekV3Bridge._build_config() expects hf_config.rope_theta directly. - if not hasattr(hf_config, "rope_theta"): - rope_params = getattr(hf_config, "rope_parameters", None) or {} - hf_config.rope_theta = rope_params.get("rope_theta", 1000000) - super().__init__(hf_config, **kwargs) - - _DSA_ATTENTION_MAPPING = { - "self_attention.wq_b.weight": ["model.layers.{layer_number}.self_attn.indexer.wq_b.weight"], - "self_attention.wk.weight": ["model.layers.{layer_number}.self_attn.indexer.wk.weight"], - "self_attention.weights_proj.weight": ["model.layers.{layer_number}.self_attn.indexer.weights_proj.weight"], - "self_attention.k_norm.weight": ["model.layers.{layer_number}.self_attn.indexer.k_norm.weight"], - "self_attention.k_norm.bias": ["model.layers.{layer_number}.self_attn.indexer.k_norm.bias"], - } - _ATTENTION_MAPPING = {**DeepseekV3Bridge._ATTENTION_MAPPING, **_DSA_ATTENTION_MAPPING} - - def _weight_to_hf_format( - self, mcore_weights_name: str, mcore_weights: torch.Tensor - ) -> tuple[list[str], list[torch.Tensor]]: - """Apply rope reordering when exporting DSA attention weights to HF format. - - Our training uses last half for rope while DeepSeek uses first half, - so we swap the two halves. - """ - if "self_attention.wq_b.weight" in mcore_weights_name: - hf_names = self._weight_name_mapping_mcore_to_hf(mcore_weights_name) - wq_b = mcore_weights - wq_b = wq_b.view(-1, 128, wq_b.shape[-1]) # hard code 128 - wq_b = torch.cat([wq_b[:, 64:], wq_b[:, :64]], dim=1).view(-1, wq_b.shape[-1]) - return hf_names, [wq_b] - elif "self_attention.wk.weight" in mcore_weights_name: - hf_names = self._weight_name_mapping_mcore_to_hf(mcore_weights_name) - wk = mcore_weights - wk = torch.cat([wk[64:], wk[:64]], dim=0) - return hf_names, [wk] - elif "self_attention.k_norm.weight" in mcore_weights_name: - hf_names = self._weight_name_mapping_mcore_to_hf(mcore_weights_name) - knorm_weight = mcore_weights - knorm_weight = torch.cat([knorm_weight[64:], knorm_weight[:64]], dim=0) - return hf_names, [knorm_weight] - elif "self_attention.k_norm.bias" in mcore_weights_name: - hf_names = self._weight_name_mapping_mcore_to_hf(mcore_weights_name) - knorm_bias = mcore_weights - knorm_bias = torch.cat([knorm_bias[64:], knorm_bias[:64]], dim=0) - return hf_names, [knorm_bias] - return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) - - def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor: - """Apply inverse rope reordering when importing DSA attention weights from HF format. - - The swap operation is its own inverse: swap the two halves back. - """ - if "self_attention.wq_b.weight" in mcore_weights_name: - wq_b = hf_weights[0] - wq_b = wq_b.view(-1, 128, wq_b.shape[-1]) # hard code 128 - wq_b = torch.cat([wq_b[:, 64:], wq_b[:, :64]], dim=1).view(-1, wq_b.shape[-1]) - return wq_b - elif "self_attention.wk.weight" in mcore_weights_name: - wk = hf_weights[0] - wk = torch.cat([wk[64:], wk[:64]], dim=0) - return wk - elif "self_attention.k_norm.weight" in mcore_weights_name: - knorm_weight = hf_weights[0] - knorm_weight = torch.cat([knorm_weight[64:], knorm_weight[:64]], dim=0) - return knorm_weight - elif "self_attention.k_norm.bias" in mcore_weights_name: - knorm_bias = hf_weights[0] - knorm_bias = torch.cat([knorm_bias[64:], knorm_bias[:64]], dim=0) - return knorm_bias - return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) diff --git a/vime_plugins/mbridge/glm4.py b/vime_plugins/mbridge/glm4.py deleted file mode 100644 index ef9e6ea70..000000000 --- a/vime_plugins/mbridge/glm4.py +++ /dev/null @@ -1,109 +0,0 @@ -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec - -from mbridge.core import LLMBridge, register_model - - -@register_model("glm4") -class GLM4Bridge(LLMBridge): - """ - Bridge implementation for Qwen2 models. - - This class extends LLMBridge to provide specific configurations and - optimizations for Qwen2 models, handling the conversion between - Hugging Face Qwen2 format and Megatron-Core. - """ - - _DIRECT_MAPPING = { - "embedding.word_embeddings.weight": "model.embed_tokens.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - "output_layer.weight": "lm_head.weight", - } - _ATTENTION_MAPPING = { - "self_attention.linear_proj.weight": ["model.layers.{layer_number}.self_attn.o_proj.weight"], - "self_attention.linear_qkv.layer_norm_weight": ["model.layers.{layer_number}.input_layernorm.weight"], - "self_attention.q_layernorm.weight": ["model.layers.{layer_number}.self_attn.q_norm.weight"], - "self_attention.k_layernorm.weight": ["model.layers.{layer_number}.self_attn.k_norm.weight"], - "self_attention.linear_qkv.weight": [ - "model.layers.{layer_number}.self_attn.q_proj.weight", - "model.layers.{layer_number}.self_attn.k_proj.weight", - "model.layers.{layer_number}.self_attn.v_proj.weight", - ], - "self_attention.linear_qkv.bias": [ - "model.layers.{layer_number}.self_attn.q_proj.bias", - "model.layers.{layer_number}.self_attn.k_proj.bias", - "model.layers.{layer_number}.self_attn.v_proj.bias", - ], - } - _MLP_MAPPING = { - "mlp.linear_fc1.weight": [ - "model.layers.{layer_number}.mlp.gate_up_proj.weight", - ], - "mlp.linear_fc1.layer_norm_weight": ["model.layers.{layer_number}.post_attention_layernorm.weight"], - "mlp.linear_fc2.weight": ["model.layers.{layer_number}.mlp.down_proj.weight"], - } - - def _build_config(self): - """ - Build the configuration for Qwen2 models. - - Configures Qwen2-specific parameters such as QKV bias settings and - layer normalization options. - - Returns: - TransformerConfig: Configuration object for Qwen2 models - """ - return self._build_base_config( - # qwen2 - add_qkv_bias=True, - qk_layernorm=False, - post_mlp_layernorm=True, - post_self_attn_layernorm=True, - rotary_interleaved=True, - ) - - def _get_transformer_layer_spec(self): - """ - Gets the transformer layer specification. - - Creates and returns a specification for the transformer layers based on - the current configuration. - - Returns: - TransformerLayerSpec: Specification for transformer layers - - Raises: - AssertionError: If normalization is not RMSNorm - """ - transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( - post_self_attn_layernorm=True, - post_mlp_layernorm=True, - ) - return transformer_layer_spec - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - """ - Map MCore weight names to Hugging Face weight names. - - Args: - mcore_weights_name: MCore weight name - - Returns: - list: Corresponding Hugging Face weight names - """ - assert "_extra_state" not in mcore_weights_name, "extra_state should not be loaded" - - if mcore_weights_name in self._DIRECT_MAPPING: - return [self._DIRECT_MAPPING[mcore_weights_name]] - - if "post_self_attn_layernorm" in mcore_weights_name: - layer_number = mcore_weights_name.split(".")[2] - return [f"model.layers.{layer_number}.post_self_attn_layernorm.weight"] - elif "post_mlp_layernorm" in mcore_weights_name: - layer_number = mcore_weights_name.split(".")[2] - return [f"model.layers.{layer_number}.post_mlp_layernorm.weight"] - elif "self_attention" in mcore_weights_name: - return self._weight_name_mapping_attention(mcore_weights_name) - elif "mlp" in mcore_weights_name: - return self._weight_name_mapping_mlp(mcore_weights_name) - else: - raise NotImplementedError(f"Unsupported parameter name: {mcore_weights_name}") diff --git a/vime_plugins/mbridge/glm4moe.py b/vime_plugins/mbridge/glm4moe.py deleted file mode 100644 index 8bfd65aa1..000000000 --- a/vime_plugins/mbridge/glm4moe.py +++ /dev/null @@ -1,122 +0,0 @@ -import re - -from mbridge.core import register_model -from mbridge.models import Qwen2Bridge, Qwen2MoEBridge - - -@register_model("glm4_moe") -class GLM4MoEBridge(Qwen2MoEBridge): - """ - Bridge implementation for Qwen2 models. - - This class extends LLMBridge to provide specific configurations and - optimizations for Qwen2 models, handling the conversion between - Hugging Face Qwen2 format and Megatron-Core. - """ - - _MLP_MAPPING = { - **(Qwen2MoEBridge._MLP_MAPPING), - **(Qwen2Bridge._MLP_MAPPING), - "mlp.router.expert_bias": ["model.layers.{layer_number}.mlp.gate.e_score_correction_bias"], - "shared_experts.linear_fc1.weight": [ - "model.layers.{layer_number}.mlp.shared_experts.gate_proj.weight", - "model.layers.{layer_number}.mlp.shared_experts.up_proj.weight", - ], - "shared_experts.linear_fc2.weight": ["model.layers.{layer_number}.mlp.shared_experts.down_proj.weight"], - } - - _MTP_MAPPING = { - "enorm.weight": ["model.layers.{layer_number}.enorm.weight"], - "hnorm.weight": ["model.layers.{layer_number}.hnorm.weight"], - "eh_proj.weight": ["model.layers.{layer_number}.eh_proj.weight"], - "final_layernorm.weight": ["model.layers.{layer_number}.shared_head.norm.weight"], - } - - def _weight_name_mapping_mtp(self, name: str, num_layers: int) -> str: - convert_names = [] - for keyword, mapping_names in self._MTP_MAPPING.items(): - if keyword in name: - convert_names.extend([x.format(layer_number=num_layers) for x in mapping_names]) - break - elif "mlp" in name: - mtp_layer_index = int(re.findall(r"mtp\.layers\.(\d+)\.", name)[0]) - name_ = re.sub( - r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name - ) - convert_names = self._weight_name_mapping_mlp(name_) - break - elif "self_attention" in name: - mtp_layer_index = int(re.findall(r"mtp\.layers.(\d+)\.", name)[0]) - name_ = re.sub( - r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name - ) - convert_names = self._weight_name_mapping_attention(name_) - break - - if len(convert_names) == 0: - raise NotImplementedError(f"Unsupported parameter name: {name}") - return convert_names - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - """ - Map MCore weight names to Hugging Face weight names. - - Args: - mcore_weights_name: MCore weight name - - Returns: - list: Corresponding Hugging Face weight names - """ - assert "_extra_state" not in mcore_weights_name, "extra_state should not be loaded" - direct_name_mapping = { - "embedding.word_embeddings.weight": "model.embed_tokens.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - "output_layer.weight": "lm_head.weight", - } - if mcore_weights_name in direct_name_mapping: - return [direct_name_mapping[mcore_weights_name]] - - if "mtp" in mcore_weights_name: # first check mtp - return self._weight_name_mapping_mtp(mcore_weights_name, self.config.num_layers) - elif "self_attention" in mcore_weights_name: - return self._weight_name_mapping_attention(mcore_weights_name) - elif "mlp" in mcore_weights_name: - return self._weight_name_mapping_mlp(mcore_weights_name) - else: - raise NotImplementedError(f"Unsupported parameter name: {mcore_weights_name}") - - def _build_config(self): - """ - Build the configuration for Qwen2 models. - - Configures Qwen2-specific parameters such as QKV bias settings and - layer normalization options. - - Returns: - TransformerConfig: Configuration object for Qwen2 models - """ - return self._build_base_config( - use_cpu_initialization=False, - # MoE specific - moe_ffn_hidden_size=self.hf_config.moe_intermediate_size, - moe_router_bias_update_rate=0.001, - moe_router_topk=self.hf_config.num_experts_per_tok, - num_moe_experts=self.hf_config.n_routed_experts, - # moe_router_load_balancing_type="aux_loss", - moe_router_load_balancing_type="none", # default None for RL - moe_grouped_gemm=True, - moe_router_score_function="sigmoid", - moe_router_enable_expert_bias=True, - moe_router_pre_softmax=True, - # Other optimizations - persist_layer_norm=True, - bias_activation_fusion=True, - bias_dropout_fusion=True, - # GLM specific - qk_layernorm=self.hf_config.use_qk_norm, - add_qkv_bias=True, - add_bias_linear=False, - # post_mlp_layernorm=True, - # post_self_attn_layernorm=True, - rotary_interleaved=True, - ) diff --git a/vime_plugins/mbridge/glm4moe_lite.py b/vime_plugins/mbridge/glm4moe_lite.py deleted file mode 100644 index ceaa8d86e..000000000 --- a/vime_plugins/mbridge/glm4moe_lite.py +++ /dev/null @@ -1,76 +0,0 @@ -import torch -from mbridge.core import register_model -from mbridge.core.safetensor_io import SafeTensorIO -from mbridge.models import DeepseekV3Bridge - - -@register_model("glm4_moe_lite") -class GLM4MoELiteBridge(DeepseekV3Bridge): - """ - Bridge for GLM-4.7-Flash (glm4_moe_lite) models. - - Extends DeepseekV3Bridge with: - - Dynamic MTP layer indexing (parent hardcodes layer 61 for DeepSeek V3) - - Standard bf16 safetensor loading (parent uses FP8 dequant for DeepSeek V3) - """ - - def __init__(self, hf_config, **kwargs): - # Patch rope_theta: GLM-4.7-Flash stores it in rope_parameters dict, - # but DeepseekV3Bridge._build_config() expects hf_config.rope_theta directly. - if not hasattr(hf_config, "rope_theta"): - rope_params = getattr(hf_config, "rope_parameters", None) or {} - hf_config.rope_theta = rope_params.get("rope_theta", 1000000) - super().__init__(hf_config, **kwargs) - # Override the shared state dict mapping with dynamic layer index. - # DeepseekV3Bridge hardcodes layer 61; GLM-4.7-Flash uses num_hidden_layers (47). - n = hf_config.num_hidden_layers - if getattr(hf_config, "num_nextn_predict_layers", 0) and n: - self._SHARED_STATE_DICT_MAPPING = { - "embedding.word_embeddings.weight": [ - "model.embed_tokens.weight", - f"model.layers.{n}.embed_tokens.weight", - ], - "output_layer.weight": [ - "lm_head.weight", - f"model.layers.{n}.shared_head.head.weight", - ], - } - - def _get_safetensor_io(self, weights_path: str): - """Use standard SafeTensorIO — GLM-4.7-Flash ships bf16 safetensors, not FP8.""" - return SafeTensorIO(self._get_actual_hf_path(weights_path)) - - def _weight_to_hf_format( - self, mcore_weights_name: str, mcore_weights: torch.Tensor - ) -> tuple[list[str], list[torch.Tensor]]: - """Handle shared embedding/output weights for MTP with dynamic layer count.""" - if ( - self.config.mtp_num_layers is not None - and self.config.mtp_num_layers >= 1 - and mcore_weights_name in self._SHARED_STATE_DICT_MAPPING - ): - hf_names = self._SHARED_STATE_DICT_MAPPING[mcore_weights_name] - return hf_names, [mcore_weights] * len(hf_names) - # Skip DeepseekV3Bridge's _weight_to_hf_format (hardcoded 61) and go to Bridge base - return super(DeepseekV3Bridge, self)._weight_to_hf_format(mcore_weights_name, mcore_weights) - - def _convert_mtp_param(self, name: str) -> list[str]: - """Convert MTP parameter names with dynamic layer count (not hardcoded 61).""" - assert self.config.mtp_num_layers == 1, "only support one mtp layer for now" - n = self.config.num_layers - direct_name_mapping = { - "mtp.layers.0.enorm.weight": f"model.layers.{n}.enorm.weight", - "mtp.layers.0.hnorm.weight": f"model.layers.{n}.hnorm.weight", - "mtp.layers.0.eh_proj.weight": f"model.layers.{n}.eh_proj.weight", - "mtp.layers.0.final_layernorm.weight": f"model.layers.{n}.shared_head.norm.weight", - } - if name in direct_name_mapping: - return [direct_name_mapping[name]] - assert "mtp.layers.0.transformer_layer" in name, f"mtp not found in {name}" - proxy_name = name.replace("mtp.layers.0.transformer_layer", f"decoder.layers.{n}") - if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: - return self._weight_name_mapping_attention(proxy_name) - elif "mlp" in proxy_name: - return self._weight_name_mapping_mlp(proxy_name) - else: - raise NotImplementedError(f"Unsupported MTP parameter name: {name}") diff --git a/vime_plugins/mbridge/gpt_oss.py b/vime_plugins/mbridge/gpt_oss.py deleted file mode 100644 index fce914d11..000000000 --- a/vime_plugins/mbridge/gpt_oss.py +++ /dev/null @@ -1,125 +0,0 @@ -from mbridge.core import register_model -from mbridge.models import Qwen2Bridge, Qwen2MoEBridge - - -@register_model("gpt_oss") -class GptOssBridge(Qwen2MoEBridge): - """ - Bridge implementation for GPT-OSS models. - - Handles weight conversion between preprocessed GPT-OSS HF format - (BF16 per-expert) and Megatron-Core. - - Key differences from Qwen2MoE: - - All layers are MoE (no dense layers, no shared expert) - - Has learnable softmax offset (sinks) - - Has attention bias (q/k/v/o_proj.bias) - - Has router bias - - Has expert bias (gate/up/down_proj.bias) - """ - - _ATTENTION_MAPPING = { - **(Qwen2Bridge._ATTENTION_MAPPING), - "self_attention.linear_proj.bias": ["model.layers.{layer_number}.self_attn.o_proj.bias"], - "self_attention.core_attention.softmax_offset": ["model.layers.{layer_number}.self_attn.sinks"], - } - - _MLP_MAPPING = { - "pre_mlp_layernorm.weight": ["model.layers.{layer_number}.post_attention_layernorm.weight"], - "mlp.router.weight": ["model.layers.{layer_number}.mlp.router.weight"], - "mlp.router.bias": ["model.layers.{layer_number}.mlp.router.bias"], - # Expert biases (must be checked before weight patterns) - "mlp.experts.linear_fc1.bias": [ - "model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.bias", - "model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.bias", - ], - "mlp.experts.linear_fc2.bias": [ - "model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.bias", - ], - # Expert weights - "mlp.experts.linear_fc1.weight": [ - "model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", - "model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", - ], - "mlp.experts.linear_fc2.weight": [ - "model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight", - ], - } - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - assert "_extra_state" not in mcore_weights_name, "extra_state should not be loaded" - - if mcore_weights_name in self._DIRECT_MAPPING: - return [self._DIRECT_MAPPING[mcore_weights_name]] - - if "self_attention" in mcore_weights_name: - return self._weight_name_mapping_attention(mcore_weights_name) - elif "mlp" in mcore_weights_name: - return self._weight_name_mapping_mlp(mcore_weights_name) - elif "pre_mlp_layernorm" in mcore_weights_name: - return self._weight_name_mapping_mlp(mcore_weights_name) - else: - raise NotImplementedError(f"Unsupported parameter name: {mcore_weights_name}") - - def _weight_name_mapping_mlp(self, name: str) -> list[str]: - """Override to handle expert bias names correctly. - - Base class extracts expert_id by splitting on 'weight', which fails - for bias parameters. We extract expert_id from after 'bias' as well. - """ - layer_number = name.split(".")[2] - convert_names = [] - for keyword, mapping_names in self._MLP_MAPPING.items(): - if keyword in name: - if "{expert_id}" in mapping_names[0]: - # Extract expert_id from end of name (after weight/bias) - if "weight" in name.split(".")[-1]: - expert_id = name.split("weight")[-1] - elif "bias" in name.split(".")[-1]: - expert_id = name.split("bias")[-1] - else: - raise ValueError(f"Cannot extract expert_id from: {name}") - convert_names.extend( - [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names] - ) - else: - convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names]) - break - if len(convert_names) == 0: - raise NotImplementedError(f"Unsupported MLP parameter name: {name}") - return convert_names - - def _build_config(self): - return self._build_base_config( - use_cpu_initialization=False, - # MoE - moe_ffn_hidden_size=self.hf_config.intermediate_size, - moe_router_topk=self.hf_config.num_experts_per_tok, - num_moe_experts=self.hf_config.num_local_experts, - moe_router_load_balancing_type="none", - moe_grouped_gemm=True, - moe_router_score_function="softmax", - moe_router_pre_softmax=False, - # GPT-OSS specific - add_qkv_bias=True, - add_bias_linear=True, - qk_layernorm=False, - persist_layer_norm=True, - bias_activation_fusion=False, - bias_dropout_fusion=False, - # SWA - window_size=(self.hf_config.sliding_window, 0), - window_attn_skip_freq=2, - # Learnable softmax - softmax_type="learnable", - # Quick GeGLU - glu_linear_offset=1.0, - activation_func_clamp_value=getattr(self.hf_config, "swiglu_limit", 7.0), - # RoPE - rotary_interleaved=False, - ) - - def _get_transformer_layer_spec(self): - from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec - - return get_gpt_layer_with_transformer_engine_spec() diff --git a/vime_plugins/mbridge/mimo.py b/vime_plugins/mbridge/mimo.py deleted file mode 100644 index a45e7cfdf..000000000 --- a/vime_plugins/mbridge/mimo.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -import torch -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec - -from mbridge.core import register_model -from mbridge.models import Qwen2Bridge - - -@register_model("mimo") -class MimoBridge(Qwen2Bridge): - """ - Bridge implementation for Mimo models. - - This class extends Qwen2Bridge to provide specific configurations and - optimizations for Mimo models, handling the conversion between - Hugging Face Mimo format and Megatron-Core. - - MiMo adds MTP (Multi-Token Prediction) layers on top of Qwen2 architecture. - """ - - def _build_config(self): - """Override to add MTP configuration.""" - hf_config = self.hf_config - - # Add MTP configuration if present - mtp_args = {} - if "num_nextn_predict_layers" in hf_config: - mtp_args["mtp_num_layers"] = hf_config.num_nextn_predict_layers - - return self._build_base_config( - add_qkv_bias=True, - qk_layernorm=False, - **mtp_args, - ) - - def _get_gptmodel_args(self) -> dict: - """Override to add MTP block spec if needed.""" - ret = super()._get_gptmodel_args() - - # Add MTP block spec if MTP layers are present - if self.config.mtp_num_layers is not None: - transformer_layer_spec = self.config - mtp_block_spec = get_gpt_mtp_block_spec(self.config, transformer_layer_spec, use_transformer_engine=True) - ret["mtp_block_spec"] = mtp_block_spec - - return ret - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - """Override to handle MTP layer mappings.""" - # Check if this is an MTP layer weight - if "mtp" in mcore_weights_name: - return self._convert_mtp_param(mcore_weights_name) - - # Otherwise use parent class mapping - return super()._weight_name_mapping_mcore_to_hf(mcore_weights_name) - - def _convert_mtp_param(self, name: str) -> list[str]: - """Convert MTP layer parameters from MCore to HF format.""" - # For now, assume single MTP layer support - if "mtp.layers." not in name: - raise NotImplementedError(f"Invalid MTP parameter name: {name}") - - # Get the MTP layer index - parts = name.split(".") - mtp_layer_idx = parts[2] # mtp.layers.{idx} - - # Direct mappings for MTP-specific components - direct_name_mapping = { - f"mtp.layers.{mtp_layer_idx}.enorm.weight": f"model.mtp_layers.{mtp_layer_idx}.token_layernorm.weight", - f"mtp.layers.{mtp_layer_idx}.hnorm.weight": f"model.mtp_layers.{mtp_layer_idx}.hidden_layernorm.weight", - f"mtp.layers.{mtp_layer_idx}.eh_proj.weight": f"model.mtp_layers.{mtp_layer_idx}.input_proj.weight", - f"mtp.layers.{mtp_layer_idx}.final_layernorm.weight": f"model.mtp_layers.{mtp_layer_idx}.final_layernorm.weight", - } - - if name in direct_name_mapping: - return [direct_name_mapping[name]] - - # Handle transformer components within MTP - # Check if this is a transformer_layer component - if "transformer_layer" in name: - # Create a proxy name to use with parent class methods - # Convert mtp.layers.{idx}.transformer_layer.* to decoder.layers.{idx}.* - proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", - f"decoder.layers.{mtp_layer_idx}", - ) - - if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: - convert_names = super()._weight_name_mapping_attention(proxy_name) - elif "mlp" in proxy_name: - convert_names = super()._weight_name_mapping_mlp(proxy_name) - else: - raise NotImplementedError(f"Unsupported transformer component in MTP: {name}") - - # Replace the layer index in converted names to point to mtp_layers - convert_names = [ - cn.replace(f"model.layers.{mtp_layer_idx}", f"model.mtp_layers.{mtp_layer_idx}") - for cn in convert_names - ] - return convert_names - else: - raise NotImplementedError(f"Unsupported MTP parameter name: {name}") - return convert_names - - def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor: - """Swap halves of eh_proj weights before handing off to Megatron-Core.""" - weight = super()._weight_to_mcore_format(mcore_weights_name, hf_weights) - if mcore_weights_name.endswith("eh_proj.weight"): - first_half, second_half = weight.chunk(2, dim=1) - weight = torch.cat([second_half, first_half], dim=1) - return weight - - def _weight_to_hf_format( - self, mcore_weights_name: str, mcore_weights: torch.Tensor - ) -> tuple[list[str], list[torch.Tensor]]: - """Swap halves back when exporting eh_proj weights to HuggingFace format.""" - if mcore_weights_name.endswith("eh_proj.weight"): - first_half, second_half = mcore_weights.chunk(2, dim=1) - mcore_weights = torch.cat([second_half, first_half], dim=1) - return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) diff --git a/vime_plugins/mbridge/minimax_m2.py b/vime_plugins/mbridge/minimax_m2.py deleted file mode 100644 index b5ade031e..000000000 --- a/vime_plugins/mbridge/minimax_m2.py +++ /dev/null @@ -1,63 +0,0 @@ -from mbridge.core import register_model -from mbridge.models import Qwen2MoEBridge - - -@register_model("minimax_m2") -class MiniMaxM2Bridge(Qwen2MoEBridge): - """ - Bridge for MiniMax-M2.5 (229B MoE). - - Key differences from standard Qwen2MoE: - - HF uses `block_sparse_moe` prefix (not `mlp`) with expert naming w1/w2/w3 - - Full-dimension QK Norm: custom SelfAttention uses `q_norm`/`k_norm` fields - (NOT the default `q_layernorm`/`k_layernorm`), so state_dict key is - `self_attention.q_norm.weight` / `self_attention.k_norm.weight` - - Sigmoid router with e_score_correction_bias - - Partial RoPE (rotary_percent=0.5) - - No shared experts - """ - - _ATTENTION_MAPPING = { - **Qwen2MoEBridge._ATTENTION_MAPPING, - # Override QK norm: custom MiniMaxM2SelfAttention uses self.q_norm / self.k_norm - # instead of the default self.q_layernorm / self.k_layernorm - "self_attention.q_norm.weight": ["model.layers.{layer_number}.self_attn.q_norm.weight"], - "self_attention.k_norm.weight": ["model.layers.{layer_number}.self_attn.k_norm.weight"], - } - - _MLP_MAPPING = { - "pre_mlp_layernorm": ["model.layers.{layer_number}.post_attention_layernorm.weight"], - "mlp.router.weight": ["model.layers.{layer_number}.block_sparse_moe.gate.weight"], - "mlp.router.expert_bias": ["model.layers.{layer_number}.block_sparse_moe.e_score_correction_bias"], - "mlp.experts.linear_fc1": [ - "model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w1.weight", # gate_proj - "model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w3.weight", # up_proj - ], - "mlp.experts.linear_fc2": [ - "model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w2.weight", # down_proj - ], - } - - def _build_config(self): - return self._build_base_config( - use_cpu_initialization=False, - persist_layer_norm=True, - bias_activation_fusion=True, - bias_dropout_fusion=True, - # MoE config - moe_ffn_hidden_size=self.hf_config.intermediate_size, - moe_router_topk=self.hf_config.num_experts_per_tok, - num_moe_experts=self.hf_config.num_local_experts, - moe_router_score_function="sigmoid", - moe_router_enable_expert_bias=True, - moe_router_pre_softmax=True, - moe_router_dtype="fp32", - moe_grouped_gemm=True, - moe_router_load_balancing_type="none", - # Attention config - qk_layernorm=True, - rotary_percent=0.5, - add_qkv_bias=False, - add_bias_linear=False, - rotary_interleaved=False, - ) diff --git a/vime_plugins/mbridge/qwen3_5.py b/vime_plugins/mbridge/qwen3_5.py deleted file mode 100644 index d01094ecc..000000000 --- a/vime_plugins/mbridge/qwen3_5.py +++ /dev/null @@ -1,355 +0,0 @@ -import inspect - -import torch -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec - -from mbridge.core import register_model -from mbridge.models import Qwen2MoEBridge - - -@register_model(["qwen3_5", "qwen3_5_moe"]) -class Qwen3_5Bridge(Qwen2MoEBridge): - """ - Bridge for Qwen3.5 models (both dense and MoE variants). - Qwen3.5 is a VLM model with weights under model.language_model.layers prefix, - separate in_proj_qkv + in_proj_z for linear attention, and nested text_config. - """ - - _DIRECT_MAPPING = { - "embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight", - "decoder.final_layernorm.weight": "model.language_model.norm.weight", - "output_layer.weight": "lm_head.weight", - } - - _ATTENTION_MAPPING = { - "self_attention.linear_proj.weight": ["model.language_model.layers.{layer_number}.self_attn.o_proj.weight"], - "self_attention.linear_qkv.layer_norm_weight": [ - "model.language_model.layers.{layer_number}.input_layernorm.weight" - ], - "self_attention.q_layernorm.weight": ["model.language_model.layers.{layer_number}.self_attn.q_norm.weight"], - "self_attention.k_layernorm.weight": ["model.language_model.layers.{layer_number}.self_attn.k_norm.weight"], - "self_attention.linear_qkv.weight": [ - "model.language_model.layers.{layer_number}.self_attn.q_proj.weight", - "model.language_model.layers.{layer_number}.self_attn.k_proj.weight", - "model.language_model.layers.{layer_number}.self_attn.v_proj.weight", - ], - "self_attention.linear_qkv.bias": [ - "model.language_model.layers.{layer_number}.self_attn.q_proj.bias", - "model.language_model.layers.{layer_number}.self_attn.k_proj.bias", - "model.language_model.layers.{layer_number}.self_attn.v_proj.bias", - ], - } | { - f"self_attention.{weight_name}": ["model.language_model.layers.{layer_number}." + weight_name] - for weight_name in [ - "input_layernorm.weight", - # linear attn - "linear_attn.A_log", - "linear_attn.conv1d.weight", - "linear_attn.dt_bias", - "linear_attn.in_proj_a.weight", - "linear_attn.in_proj_b.weight", - "linear_attn.in_proj_qkv.weight", - "linear_attn.in_proj_z.weight", - "linear_attn.norm.weight", - "linear_attn.out_proj.weight", - # gated attn (full attention layers) - "self_attn.k_norm.weight", - "self_attn.k_proj.weight", - "self_attn.o_proj.weight", - "self_attn.q_norm.weight", - "self_attn.q_proj.weight", - "self_attn.v_proj.weight", - ] - } - - _MLP_MAPPING = { - "mlp.linear_fc1.weight": [ - "model.language_model.layers.{layer_number}.mlp.gate_proj.weight", - "model.language_model.layers.{layer_number}.mlp.up_proj.weight", - ], - "mlp.linear_fc1.layer_norm_weight": [ - "model.language_model.layers.{layer_number}.post_attention_layernorm.weight" - ], - "mlp.linear_fc2.weight": ["model.language_model.layers.{layer_number}.mlp.down_proj.weight"], - # MoE mappings - "shared_experts.linear_fc1.weight": [ - "model.language_model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight", - "model.language_model.layers.{layer_number}.mlp.shared_expert.up_proj.weight", - ], - "pre_mlp_layernorm": ["model.language_model.layers.{layer_number}.post_attention_layernorm.weight"], - "shared_experts.linear_fc2.weight": [ - "model.language_model.layers.{layer_number}.mlp.shared_expert.down_proj.weight" - ], - "mlp.router.weight": ["model.language_model.layers.{layer_number}.mlp.gate.weight"], - "shared_experts.gate_weight": ["model.language_model.layers.{layer_number}.mlp.shared_expert_gate.weight"], - # Fused expert format: single 3D tensor for all experts - "mlp.experts.linear_fc1": [ - "model.language_model.layers.{layer_number}.mlp.experts.gate_up_proj", - ], - "mlp.experts.linear_fc2": ["model.language_model.layers.{layer_number}.mlp.experts.down_proj"], - } - - # MTP layer uses individual expert format (not fused) - _MTP_MLP_MAPPING = { - "mlp.experts.linear_fc1": [ - "mtp.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", - "mtp.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", - ], - "mlp.experts.linear_fc2": ["mtp.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight"], - } - - # Override to make ffn_hidden_size optional (Qwen3.5 MoE has no intermediate_size) - _CONFIG_MAPPING = { - "num_layers": "num_hidden_layers", - "hidden_size": "hidden_size", - "num_attention_heads": "num_attention_heads", - "num_query_groups": "num_key_value_heads", - "ffn_hidden_size": ("intermediate_size", None), - "attention_dropout": "attention_dropout", - "layernorm_epsilon": "rms_norm_eps", - "hidden_dropout": ("hidden_dropout", 0.0), - "kv_channels": ("head_dim", None), - } - - def _get_text_config(self): - """Get the text config, handling VLM nesting.""" - if hasattr(self.hf_config, "text_config"): - return self.hf_config.text_config - return self.hf_config - - def _adjust_mapping_for_shared_weights(self): - text_config = self._get_text_config() - tie_word_embeddings = getattr(text_config, "tie_word_embeddings", False) or getattr( - self.hf_config, "tie_word_embeddings", False - ) - if tie_word_embeddings: - self._DIRECT_MAPPING = dict(self._DIRECT_MAPPING) - self._DIRECT_MAPPING["output_layer.weight"] = "model.language_model.embed_tokens.weight" - - def _supports_transformer_config_kwarg(self, kwarg_name: str) -> bool: - """Check whether the current TransformerConfig accepts a given kwarg.""" - transformer_config_class = getattr(self, "TransformerConfigClass", None) - if transformer_config_class is None: - return True - - dataclass_fields = getattr(transformer_config_class, "__dataclass_fields__", None) - if dataclass_fields is not None: - return kwarg_name in dataclass_fields - - try: - signature = inspect.signature(transformer_config_class) - except (TypeError, ValueError): - return True - return kwarg_name in signature.parameters - - def _get_transformer_layer_spec(self, vp_stage=None): - transformer_layer_spec = super()._get_transformer_layer_spec(vp_stage) - self._last_transformer_layer_spec = transformer_layer_spec - return transformer_layer_spec - - def _get_gptmodel_args(self) -> dict: - """Override to add MTP block spec if needed.""" - ret = super()._get_gptmodel_args() - text_config = self._get_text_config() - if getattr(text_config, "mtp_num_hidden_layers", None) is not None: - transformer_layer_spec = getattr(self, "_last_transformer_layer_spec", None) - if transformer_layer_spec is None: - transformer_layer_spec = self._get_transformer_layer_spec() - mtp_block_spec = get_gpt_mtp_block_spec(self.config, transformer_layer_spec, use_transformer_engine=True) - ret["mtp_block_spec"] = mtp_block_spec - return ret - - def _weight_name_mapping_mlp(self, name: str) -> list[str]: - """Override to handle fused expert weights. - For regular layers: experts use fused 3D format (all experts in one tensor). - For MTP layers: experts use individual format (per-expert tensors). - """ - layer_number = name.split(".")[2] - convert_names = [] - for keyword, mapping_names in self._MLP_MAPPING.items(): - if keyword in name: - if "{expert_id}" in mapping_names[0]: - expert_id = name.split("weight")[-1] - convert_names.extend( - [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names] - ) - else: - convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names]) - break - if len(convert_names) == 0: - raise NotImplementedError(f"Unsupported parameter name: {name}") - return convert_names - - def _weight_name_mapping_mtp_mlp(self, name: str) -> list[str]: - """Handle MTP MLP mappings, keeping per-expert tensors unfused for MoE layers.""" - layer_number = name.split(".")[2] - mapping = self._MTP_MLP_MAPPING if "mlp.experts.linear_fc" in name else self._MLP_MAPPING - convert_names = [] - for keyword, mapping_names in mapping.items(): - if keyword in name: - if "{expert_id}" in mapping_names[0]: - expert_id = name.split("weight")[-1] - convert_names.extend( - [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names] - ) - else: - convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names]) - break - if len(convert_names) == 0: - raise NotImplementedError(f"Unsupported parameter name: {name}") - return convert_names - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - """Override to handle MTP layer mappings.""" - if "mtp" in mcore_weights_name: - return self._convert_mtp_param(mcore_weights_name) - return super()._weight_name_mapping_mcore_to_hf(mcore_weights_name) - - def _convert_mtp_param(self, name: str) -> list[str]: - """Convert MTP layer parameters from MCore to HF format.""" - if "mtp.layers." not in name: - raise NotImplementedError(f"Invalid MTP parameter name: {name}") - - parts = name.split(".") - mtp_layer_idx = parts[2] # mtp.layers.{idx} - - direct_name_mapping = { - f"mtp.layers.{mtp_layer_idx}.eh_proj.weight": "mtp.fc.weight", - f"mtp.layers.{mtp_layer_idx}.enorm.weight": "mtp.pre_fc_norm_embedding.weight", - f"mtp.layers.{mtp_layer_idx}.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", - f"mtp.layers.{mtp_layer_idx}.final_layernorm.weight": "mtp.norm.weight", - } - - if name in direct_name_mapping: - return [direct_name_mapping[name]] - - if "transformer_layer" in name: - proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", - f"decoder.layers.{mtp_layer_idx}", - ) - - if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: - convert_names = super()._weight_name_mapping_attention(proxy_name) - elif "mlp" in proxy_name or "pre_mlp_layernorm" in proxy_name: - convert_names = self._weight_name_mapping_mtp_mlp(proxy_name) - else: - raise NotImplementedError(f"Unsupported transformer component in MTP: {name}") - - # MTP weights use model.language_model prefix in regular layers, - # but mtp.layers.{idx} directly for MTP layers - convert_names = [ - cn.replace(f"model.language_model.layers.{mtp_layer_idx}", f"mtp.layers.{mtp_layer_idx}") - for cn in convert_names - ] - return convert_names - - raise NotImplementedError(f"Unsupported MTP parameter name: {name}") - - def _weight_to_mcore_format( - self, mcore_weights_name: str, hf_weights: list[torch.Tensor] - ) -> tuple[list[str], list[torch.Tensor]]: - if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: - # merge qkv - assert len(hf_weights) == 3 - text_config = self._get_text_config() - num_key_value_heads = text_config.num_key_value_heads - hidden_dim = text_config.hidden_size - num_attention_heads = text_config.num_attention_heads - num_querys_per_group = num_attention_heads // text_config.num_key_value_heads - head_dim = getattr(text_config, "head_dim", hidden_dim // num_attention_heads) - group_dim = head_dim * num_attention_heads // num_key_value_heads - q, k, v = hf_weights - # q k v might be tp split - real_num_key_value_heads = q.shape[0] // (2 * group_dim) - q = ( - q.view( - [ - real_num_key_value_heads, - num_querys_per_group, - 2, - head_dim, - -1, - ] - ) - .transpose(1, 2) - .flatten(1, 3) - ) - k = k.view([real_num_key_value_heads, head_dim, -1]) - v = v.view([real_num_key_value_heads, head_dim, -1]) - out_shape = [-1, hidden_dim] if ".bias" not in mcore_weights_name else [-1] - - qgkv = torch.cat([q, k, v], dim=1).view(*out_shape).contiguous() - return qgkv - - # Handle fused expert weights: extract single expert from 3D fused tensor - if "mlp.experts.linear_fc" in mcore_weights_name and len(hf_weights) == 1: - w = hf_weights[0] - if w.dim() == 3: - # Extract local expert_id from name like "...linear_fc1.weight42" - local_expert_id = int(mcore_weights_name.split("weight")[-1]) - # When using Expert Parallelism (EP), the local expert_id is relative - # to this EP rank. We need to convert to global expert_id to index - # into the full HF fused tensor [num_experts, ...]. - from megatron.core import mpu - - ep_size = mpu.get_expert_model_parallel_world_size() - if ep_size > 1: - ep_rank = mpu.get_expert_model_parallel_rank() - num_local_experts = w.shape[0] // ep_size - global_expert_id = ep_rank * num_local_experts + local_expert_id - else: - global_expert_id = local_expert_id - expert_w = w[global_expert_id] # (out_features, in_features) - return expert_w.contiguous() - - return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) - - def _weight_to_hf_format( - self, mcore_weights_name: str, mcore_weights: torch.Tensor - ) -> tuple[list[str], list[torch.Tensor]]: - return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) - - def _build_config(self): - text_config = self._get_text_config() - - mtp_args = {} - if hasattr(text_config, "mtp_num_hidden_layers"): - mtp_args["mtp_num_layers"] = text_config.mtp_num_hidden_layers - - base_kwargs = dict( - text_config_key="text_config" if hasattr(self.hf_config, "text_config") else None, - use_cpu_initialization=False, - # Other optimizations - persist_layer_norm=True, - bias_activation_fusion=True, - bias_dropout_fusion=True, - # Qwen3.5 specific - moe_router_pre_softmax=False, - qk_layernorm=True, - attention_output_gate=True, - **mtp_args, - ) - - if self._supports_transformer_config_kwarg("use_gated_attention"): - base_kwargs["use_gated_attention"] = True - - # Handle MoE-specific config - if hasattr(text_config, "num_experts"): - base_kwargs.update( - moe_ffn_hidden_size=text_config.moe_intermediate_size, - moe_shared_expert_intermediate_size=getattr(text_config, "shared_expert_intermediate_size", None), - moe_router_bias_update_rate=0.001, - moe_router_topk=text_config.num_experts_per_tok, - num_moe_experts=text_config.num_experts, - moe_aux_loss_coeff=text_config.router_aux_loss_coef, - moe_router_load_balancing_type="none", - moe_grouped_gemm=True, - moe_router_score_function="softmax", - moe_shared_expert_gate=True, - ) - # For MoE models without intermediate_size, use shared_expert_intermediate_size - if not hasattr(text_config, "intermediate_size"): - base_kwargs["ffn_hidden_size"] = text_config.shared_expert_intermediate_size - - return self._build_base_config(**base_kwargs) diff --git a/vime_plugins/mbridge/qwen3_next.py b/vime_plugins/mbridge/qwen3_next.py deleted file mode 100644 index 8fd188f5d..000000000 --- a/vime_plugins/mbridge/qwen3_next.py +++ /dev/null @@ -1,173 +0,0 @@ -import torch -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec - -from mbridge.core import register_model -from mbridge.models import Qwen2MoEBridge - - -@register_model("qwen3_next") -class Qwen3NextBridge(Qwen2MoEBridge): - _ATTENTION_MAPPING = ( - Qwen2MoEBridge._ATTENTION_MAPPING - | { - f"self_attention.{weight_name}": ["model.layers.{layer_number}." + weight_name] - for weight_name in [ - "input_layernorm.weight", - # linear attn - "linear_attn.A_log", - "linear_attn.conv1d.weight", - "linear_attn.dt_bias", - "linear_attn.in_proj_ba.weight", - "linear_attn.in_proj_qkvz.weight", - "linear_attn.norm.weight", - "linear_attn.out_proj.weight", - # gated attn - "self_attn.k_norm.weight", - "self_attn.k_proj.weight", - "self_attn.o_proj.weight", - "self_attn.q_norm.weight", - "self_attn.q_proj.weight", - "self_attn.v_proj.weight", - ] - } - | { - "self_attention.linear_qkv.layer_norm_weight": ["model.layers.{layer_number}.input_layernorm.weight"], - "self_attention.linear_qkv.weight": [ - "model.layers.{layer_number}.self_attn.q_proj.weight", - "model.layers.{layer_number}.self_attn.k_proj.weight", - "model.layers.{layer_number}.self_attn.v_proj.weight", - ], - } - ) - - def _get_gptmodel_args(self) -> dict: - """Override to add MTP block spec if needed.""" - ret = super()._get_gptmodel_args() - if getattr(self.config, "mtp_num_layers", None) is not None: - transformer_layer_spec = self.config - mtp_block_spec = get_gpt_mtp_block_spec(self.config, transformer_layer_spec, use_transformer_engine=True) - ret["mtp_block_spec"] = mtp_block_spec - return ret - - def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]: - """Override to handle MTP layer mappings.""" - if "mtp" in mcore_weights_name: - return self._convert_mtp_param(mcore_weights_name) - return super()._weight_name_mapping_mcore_to_hf(mcore_weights_name) - - def _convert_mtp_param(self, name: str) -> list[str]: - """Convert MTP layer parameters from MCore to HF format.""" - if "mtp.layers." not in name: - raise NotImplementedError(f"Invalid MTP parameter name: {name}") - - parts = name.split(".") - mtp_layer_idx = parts[2] # mtp.layers.{idx} - - direct_name_mapping = { - f"mtp.layers.{mtp_layer_idx}.eh_proj.weight": "mtp.fc.weight", - f"mtp.layers.{mtp_layer_idx}.enorm.weight": "mtp.pre_fc_norm_embedding.weight", - f"mtp.layers.{mtp_layer_idx}.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", - f"mtp.layers.{mtp_layer_idx}.final_layernorm.weight": "mtp.norm.weight", - } - - if name in direct_name_mapping: - return [direct_name_mapping[name]] - - if "transformer_layer" in name: - proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", - f"decoder.layers.{mtp_layer_idx}", - ) - - if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: - convert_names = super()._weight_name_mapping_attention(proxy_name) - elif "mlp" in proxy_name or "pre_mlp_layernorm" in proxy_name: - convert_names = super()._weight_name_mapping_mlp(proxy_name) - else: - raise NotImplementedError(f"Unsupported transformer component in MTP: {name}") - - convert_names = [ - cn.replace(f"model.layers.{mtp_layer_idx}", f"mtp.layers.{mtp_layer_idx}") for cn in convert_names - ] - return convert_names - - raise NotImplementedError(f"Unsupported MTP parameter name: {name}") - - def _weight_to_mcore_format( - self, mcore_weights_name: str, hf_weights: list[torch.Tensor] - ) -> tuple[list[str], list[torch.Tensor]]: - if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: - # merge qkv - assert len(hf_weights) == 3 - num_key_value_heads = self.hf_config.num_key_value_heads - hidden_dim = self.hf_config.hidden_size - num_attention_heads = self.hf_config.num_attention_heads - num_querys_per_group = num_attention_heads // self.hf_config.num_key_value_heads - head_dim = getattr(self.hf_config, "head_dim", hidden_dim // num_attention_heads) - group_dim = head_dim * num_attention_heads // num_key_value_heads - q, k, v = hf_weights - # q k v might be tp split - real_num_key_value_heads = q.shape[0] // (2 * group_dim) - q = ( - q.view( - [ - real_num_key_value_heads, - num_querys_per_group, - 2, - head_dim, - -1, - ] - ) - .transpose(1, 2) - .flatten(1, 3) - ) - k = k.view([real_num_key_value_heads, head_dim, -1]) - v = v.view([real_num_key_value_heads, head_dim, -1]) - out_shape = [-1, hidden_dim] if ".bias" not in mcore_weights_name else [-1] - - qgkv = torch.cat([q, k, v], dim=1).view(*out_shape).contiguous() - return qgkv - - weight = super()._weight_to_mcore_format(mcore_weights_name, hf_weights) - if mcore_weights_name.endswith("eh_proj.weight"): - first_half, second_half = weight.chunk(2, dim=1) - weight = torch.cat([second_half, first_half], dim=1) - return weight - - def _weight_to_hf_format( - self, mcore_weights_name: str, mcore_weights: torch.Tensor - ) -> tuple[list[str], list[torch.Tensor]]: - if mcore_weights_name.endswith("eh_proj.weight"): - first_half, second_half = mcore_weights.chunk(2, dim=1) - mcore_weights = torch.cat([second_half, first_half], dim=1) - return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) - - def _build_config(self): - mtp_args = {} - if hasattr(self.hf_config, "num_nextn_predict_layers"): - mtp_args["mtp_num_layers"] = self.hf_config.num_nextn_predict_layers - - return self._build_base_config( - use_cpu_initialization=False, - # MoE specific - moe_ffn_hidden_size=self.hf_config.moe_intermediate_size, - moe_router_bias_update_rate=0.001, - moe_router_topk=self.hf_config.num_experts_per_tok, - num_moe_experts=self.hf_config.num_experts, - moe_aux_loss_coeff=self.hf_config.router_aux_loss_coef, - # moe_router_load_balancing_type="aux_loss", - moe_router_load_balancing_type="none", # default None for RL - moe_grouped_gemm=True, - moe_router_score_function="softmax", - # Other optimizations - persist_layer_norm=True, - bias_activation_fusion=True, - bias_dropout_fusion=True, - # Qwen specific - moe_router_pre_softmax=False, - qk_layernorm=True, - # Qwen3 Next specific - attention_output_gate=True, - moe_shared_expert_gate=True, - **mtp_args, - ) diff --git a/vime_plugins/megatron_bridge/__init__.py b/vime_plugins/megatron_bridge/__init__.py deleted file mode 100644 index b097a794d..000000000 --- a/vime_plugins/megatron_bridge/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -import vime_plugins.megatron_bridge.glm4moe_lite # noqa: F401 # register GLM-4.7-Flash bridge -import vime_plugins.megatron_bridge.glm4v_moe # noqa: F401 # register GLM-4.6V bridge diff --git a/vime_plugins/megatron_bridge/glm4moe_lite.py b/vime_plugins/megatron_bridge/glm4moe_lite.py deleted file mode 100644 index c2be2cc66..000000000 --- a/vime_plugins/megatron_bridge/glm4moe_lite.py +++ /dev/null @@ -1,428 +0,0 @@ -"""GLM-4.7-Flash (``glm4_moe_lite``) bridge for megatron.bridge. - -Registers ``Glm4MoeLiteForCausalLM`` so that ``AutoBridge.from_hf_pretrained`` -recognises GLM-4.7-Flash checkpoints and can provide a Megatron-compatible model + -weight mappings, with Multi-Token Prediction (MTP) support on the Ascend 910B NPU. - -Architecture: - MLA (Multi-Head Latent Attention, DeepSeek-V3-style) - + GLM-style MoE (64 routed experts + 1 shared expert, sigmoid router with - expert bias), subclassing ``GLM45Bridge`` to reuse its MTP plumbing. -""" - -import logging -from functools import partial - -import torch -from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge -from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping -from megatron.bridge.models.glm.glm45_bridge import GLM45Bridge -from megatron.bridge.models.glm.glm_moe_mappings import GLMExpertDownProjMapping, GLMExpertGateUpProjMapping -from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM -from megatron.bridge.models.mla_provider import MLAModelProvider -from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec -from transformers import Glm4MoeLiteForCausalLM - -try: - import transformer_engine # noqa: F401 - - HAVE_TE = True -except (ImportError, ModuleNotFoundError): - HAVE_TE = False - - -logger = logging.getLogger(__name__) - - -@MegatronModelBridge.register_bridge( - source=Glm4MoeLiteForCausalLM, - target=GPTModel, - model_type="glm4_moe_lite", -) -class GLM47MTPBridge(GLM45Bridge): - """Megatron bridge for GLM-4.7-Flash (glm4_moe_lite) with MTP support. - - GLM-4.7-Flash is an MLA model (``q_lora_rank`` in its config), so it needs - the MLAModelProvider and the MLA weight mappings. GLM45Bridge's stock - provider_bridge / mapping_registry only handle non-MLA GLM-4.5 (fused QKV); - we override both with the MLA-aware versions. Everything else - (``build_conversion_tasks``, MTP loop, fused-expert handling) is inherited. - """ - - def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): - """Convert HuggingFace config to MLAModelProvider.""" - provider_kwargs = self.hf_config_to_provider_kwargs(hf_pretrained.config) - mla_rope = provider_kwargs.pop("_mla_rope_params", None) - provider_class = self.PROVIDER_CLASS if self.PROVIDER_CLASS is not None else MLAModelProvider - provider = provider_class(**provider_kwargs) - - # Set rope type - hf_rope_scaling = getattr(hf_pretrained.config, "rope_scaling", None) - rope_type = None - if hf_rope_scaling: - rope_type = hf_rope_scaling.get("type") or hf_rope_scaling.get("rope_type") - if rope_type != "yarn": - provider.position_embedding_type = "rope" - - # Match vLLM defaults (no scaling, mscale=1.0) when HF config has no explicit rope params. - if not mla_rope: - mla_rope = {"rotary_scaling_factor": 1.0, "mscale_all_dim": 1.0} - - if mla_rope: - for key, value in mla_rope.items(): - setattr(provider, key, value) - hf_config = hf_pretrained.config - - # Use decoder block spec to properly handle moe_layer_freq (mixed dense/MoE layers) - provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE) - provider.normalization = "RMSNorm" - provider.gated_linear_unit = True - provider.add_bias_linear = False - provider.share_embeddings_and_output_weights = False - provider.multi_latent_attention = True - provider.qk_layernorm = True - - provider.moe_shared_expert_overlap = True - provider.moe_token_dispatcher_type = "alltoall" - provider.moe_router_load_balancing_type = "seq_aux_loss" - provider.moe_router_pre_softmax = True - provider.moe_grouped_gemm = True - provider.moe_router_score_function = "sigmoid" - provider.moe_permute_fusion = True - provider.moe_router_enable_expert_bias = True - provider.moe_router_dtype = "fp32" - provider.moe_router_bias_update_rate = 0 - provider.moe_aux_loss_coeff = 0.001 - - provider.persist_layer_norm = True - provider.bias_activation_fusion = True - provider.bias_dropout_fusion = True - provider.hidden_dropout = 0.0 - provider.autocast_dtype = torch.bfloat16 - provider.mtp_loss_scaling_factor = 0.3 - provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size * int( - getattr(hf_config, "n_shared_experts", 1) - ) - - provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * ( - hf_config.num_hidden_layers - hf_config.first_k_dense_replace - ) - - return provider - - def _glm_hf_config(self): - """Return the HF config across bridge revisions. - - Newer bridge revisions stash the config on ``self._hf_config`` inside - ``build_conversion_tasks``; older revisions only expose the - ``self.hf_config`` property (set during ``build_conversion_tasks`` on - the base class). The peft adapter path calls ``mapping_registry`` - *before* either attribute is populated, so guard both. - """ - hf_config = getattr(self, "_hf_config", None) - if hf_config is None: - hf_config = getattr(self, "hf_config", None) - return hf_config - - def _glm_hf_keys(self): - """Return the HF state keys in a revision-safe way. - - Newer bridge revisions cache keys on ``self._hf_keys``; older ones - reach them through ``self.hf_pretrained.state.source``. The latter - raises ``AttributeError`` on config-only paths (peft adapter export), - so fall back to ``None`` instead of propagating. - """ - hf_keys = getattr(self, "_hf_keys", None) - if hf_keys: - return hf_keys - try: - source = self.hf_pretrained.state.source - except AttributeError: - return None - return list(source.get_all_keys()) if source is not None else None - - def _uses_fused_experts(self) -> bool: - """Determine whether expert weights are fused (gate_up_proj/down_proj). - - Detection follows the base bridge's key inspection when the HF state - is available. GLM-4.7-Flash HuggingFace checkpoints ship per-expert - weights (``experts..gate_proj`` / ``up_proj`` / ``down_proj``), - not fused tensors, so the config-only fallback returns ``False``. - """ - hf_keys = self._glm_hf_keys() - if hf_keys: - if any("mlp.experts.gate_up_proj" in key for key in hf_keys) or any( - "mlp.experts.down_proj" in key for key in hf_keys - ): - return True - return False - # Config-only path: GLM-4.7-Flash uses per-expert (non-fused) weights. - return False - - def _hf_expert_suffix(self, base_name: str) -> str: - """Resolve the expert tensor suffix (``.weight`` or ``""``) safely.""" - hf_keys = self._glm_hf_keys() - if hf_keys: - if any(f"{base_name}.weight" in key for key in hf_keys): - return ".weight" - return "" - # Config-only path: GLM fused expert tensors have no .weight suffix. - return "" - - def mapping_registry(self) -> MegatronMappingRegistry: - mapping_list = [] - use_fused_experts = self._uses_fused_experts() - gate_up_suffix = self._hf_expert_suffix("mlp.experts.gate_up_proj") - down_suffix = self._hf_expert_suffix("mlp.experts.down_proj") - - param_mappings = { - # Embed - "embedding.word_embeddings.weight": "model.embed_tokens.weight", - # LM Head - "decoder.final_layernorm.weight": "model.norm.weight", - "output_layer.weight": "lm_head.weight", - } - - layer_specific_mappings = { - # Attention shared by all GLM variants - "decoder.layers.*.input_layernorm.weight": "model.layers.*.input_layernorm.weight", - "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", - "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_a_layernorm.weight", - "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", - # MLA-specific layernorm - "decoder.layers.*.self_attention.kv_layernorm.weight": "model.layers.*.self_attn.kv_a_layernorm.weight", - # MLP - "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", - "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.shared_experts.down_proj.weight", - "decoder.layers.*.mlp.shared_experts.router.weight": "model.layers.*.mlp.shared_experts.gate.weight", - "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight", - "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias", - } - - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - for megatron_param, hf_param in layer_specific_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - # Add special mappings that require parameter concatenation/transformation - mapping_list.extend( - [ - # MLA attention: individual Q/KV down/up projections (for GLM-4.7-Flash) - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_q_down_proj.weight", - hf_param="model.layers.*.self_attn.q_a_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_q_up_proj.weight", - hf_param="model.layers.*.self_attn.q_b_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_kv_down_proj.weight", - hf_param="model.layers.*.self_attn.kv_a_proj_with_mqa.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_kv_up_proj.weight", - hf_param="model.layers.*.self_attn.kv_b_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_q_up_proj.layer_norm_weight", - hf_param="model.layers.*.self_attn.q_a_layernorm.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.self_attention.linear_kv_up_proj.layer_norm_weight", - hf_param="model.layers.*.self_attn.kv_a_layernorm.weight", - ), - # Gated MLP: Combine gate and up projection matrices into single FC1 matrix - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.linear_fc1.weight", - gate="model.layers.*.mlp.gate_proj.weight", - up="model.layers.*.mlp.up_proj.weight", - ), - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", - gate="model.layers.*.mlp.shared_experts.gate_proj.weight", - up="model.layers.*.mlp.shared_experts.up_proj.weight", - ), - ] - ) - if use_fused_experts: - mapping_list.extend( - [ - GLMExpertGateUpProjMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", - hf_param=f"model.layers.*.mlp.experts.gate_up_proj{gate_up_suffix}", - ), - GLMExpertDownProjMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", - hf_param=f"model.layers.*.mlp.experts.down_proj{down_suffix}", - ), - ] - ) - else: - mapping_list.extend( - [ - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", - gate="model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", - hf_param="model.layers.*.mlp.experts.*.down_proj.weight", - ), - ] - ) - # optionally add MTP mappings - hf_config = self._glm_hf_config() - if hf_config is None: - logger.warning("No HF config found, skipping MTP mappings.") - return MegatronMappingRegistry(*mapping_list) - num_mtp_layers = getattr(hf_config, "num_nextn_predict_layers", 0) - num_transformer_layers = hf_config.num_hidden_layers - for mtp_layer in range(num_mtp_layers): - for megatron_param, hf_param in layer_specific_mappings.items(): - megatron_param = ( - megatron_param.replace(".*", ".*.transformer_layer") - .replace("decoder", "mtp") - .replace(".*", f".{mtp_layer}") - ) - hf_param = hf_param.replace("layers.*", f"layers.{mtp_layer + num_transformer_layers}") - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - # MTP specific mappings - mapping_list.extend( - [ - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.enorm.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.enorm.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.hnorm.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.hnorm.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.eh_proj.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.eh_proj.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.final_layernorm.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.shared_head.norm.weight", - ), - ] - ) - # MTP transformer layer reuses the last normal layer spec (MLA), so map - # the individual Q/KV down/up projections instead of a fused QKV. - mapping_list.extend( - [ - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_down_proj.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_a_proj.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_up_proj.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_b_proj.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_down_proj.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_a_proj_with_mqa.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_up_proj.weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_b_proj.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_q_up_proj.layer_norm_weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.q_a_layernorm.weight", - ), - AutoMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_kv_up_proj.layer_norm_weight", - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.self_attn.kv_a_layernorm.weight", - ), - ] - ) - # MTP transformer layer MLP mappings - mapping_list.extend( - [ - GatedMLPMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.mlp.linear_fc1.weight", - gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.linear_fc1.gate.weight", - up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.linear_fc1.up.weight", - ), - GatedMLPMapping( - megatron_param=f"mtp.layers.{mtp_layer}.transformer_layer.mlp.shared_experts.linear_fc1.weight", - gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.shared_experts.gate_proj.weight", - up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.shared_experts.up_proj.weight", - ), - ] - ) - if use_fused_experts: - mapping_list.extend( - [ - GLMExpertGateUpProjMapping( - megatron_param=( - f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc1.weight*" - ), - hf_param=( - f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.gate_up_proj" - f"{gate_up_suffix}" - ), - ), - GLMExpertDownProjMapping( - megatron_param=( - f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc2.weight*" - ), - hf_param=( - f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.down_proj{down_suffix}" - ), - ), - ] - ) - else: - mapping_list.extend( - [ - GatedMLPMapping( - megatron_param=( - f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc1.weight*" - ), - gate=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.gate_proj.weight", - up=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param=( - f"mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc2.weight*" - ), - hf_param=f"model.layers.{mtp_layer + num_transformer_layers}.mlp.experts.*.down_proj.weight", - ), - ] - ) - - return MegatronMappingRegistry(*mapping_list) - - -def _register_mindspeed_te_module_types(): - """Register MindSpeed TE module types for weight-mapping parallelism detection.""" - try: - from megatron.bridge.models.conversion.param_mapping import AutoMapping - except ImportError: - return - - for module_name, parallelism_type in { - "MindSpeedTEColumnParallelLinear": "column", - "MindSpeedTELayerNormColumnParallelLinear": "column", - "MindSpeedTEColumnParallelGroupedLinear": "column", - "MindSpeedTEGroupedLinear": "column", - "MindSpeedTEGroupedLinearGMM": "column", - "MindSpeedTEDotProductAttention": "column", - "MindSpeedTERowParallelGroupedLinear": "row", - "MindSpeedTELayernorm": "replicated", - "MindSpeedTELinear": "replicated", - }.items(): - AutoMapping.register_module_type(module_name, parallelism_type) - - -_register_mindspeed_te_module_types() diff --git a/vime_plugins/megatron_bridge/glm4v_moe.py b/vime_plugins/megatron_bridge/glm4v_moe.py deleted file mode 100644 index a639b905f..000000000 --- a/vime_plugins/megatron_bridge/glm4v_moe.py +++ /dev/null @@ -1,717 +0,0 @@ -""" -GLM-4.6V (glm4v_moe) bridge for megatron.bridge. - -Registers `Glm4vMoeForConditionalGeneration` so that `AutoBridge.from_hf_pretrained` -recognises GLM-4.6V checkpoints and can provide a Megatron-compatible VL model + -weight mappings. - -Architecture: - HF vision encoder (Glm4vMoeVisionModel, replicated on first PP stage) - + Megatron GPTModel (MoE language model, standard M-RoPE) -""" - -from __future__ import annotations - -import itertools -import logging -from copy import deepcopy -from dataclasses import dataclass, field - -import torch -from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge -from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping, QKVMapping, ReplicatedMapping -from megatron.bridge.models.gpt_provider import GPTModelProvider -from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync -from megatron.core import parallel_state, tensor_parallel -from megatron.core.models.gpt import GPTModel as MCoreGPTModel -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.transformer.module import MegatronModule - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# THD ↔ BSHD helpers (cf. Qwen3VL bridge) -# --------------------------------------------------------------------------- -def _thd_to_bshd(packed: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: - """Unpack THD-format [1, T, ...] to BSHD [bs, max_seq, ...] using cu_seqlens.""" - seqlens = cu_seqlens[1:] - cu_seqlens[:-1] - max_seq = seqlens.max().item() - bs = len(cu_seqlens) - 1 - out = packed.new_zeros(bs, max_seq, *packed.shape[2:]) - for i, sl in enumerate(seqlens): - out[i, :sl] = packed[0, cu_seqlens[i] : cu_seqlens[i] + sl] - return out - - -def _bshd_to_thd(unpacked: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: - """Pack BSHD [bs, max_seq, ...] back to THD [1, T, ...].""" - seqlens = cu_seqlens[1:] - cu_seqlens[:-1] - total = cu_seqlens[-1].item() - out = unpacked.new_zeros(1, total, *unpacked.shape[2:]) - for i, sl in enumerate(seqlens): - out[0, cu_seqlens[i] : cu_seqlens[i] + sl] = unpacked[i, :sl] - return out - - -def _gather_input_ids_from_cp( - input_ids: torch.Tensor, - cu_seqlens: torch.Tensor, -) -> torch.Tensor: - """Reconstruct full (global) input_ids from zigzag CP chunks. - - With zigzag CP, each CP rank r holds chunks [r] and [2*cp_size-1-r] for - each sequence. This function all-gathers across CP ranks and reassembles - the original token order so that position-ID computation sees the full - sequence. - - Args: - input_ids: Local input_ids in THD format [1, T_local]. - cu_seqlens: **Global** cumulative sequence lengths. - - Returns: - Full input_ids in THD format [1, T_global]. - """ - cp_size = parallel_state.get_context_parallel_world_size() - if cp_size <= 1: - return input_ids - - # all-gather local input_ids across CP ranks → list of [1, T_local] per rank - gathered = torch.distributed.nn.all_gather( - input_ids, group=parallel_state.get_context_parallel_group() - ) # list of cp_size tensors, each [1, T_local] - - local_cu_seqlens = cu_seqlens // cp_size - num_seqs = len(cu_seqlens) - 1 - whole_list = [] - for i in range(num_seqs): - seqlen = (cu_seqlens[i + 1] - cu_seqlens[i]).item() - chunk_size = seqlen // 2 // cp_size - # First half: rank 0 chunk, rank 1 chunk, ..., rank cp_size-1 chunk - whole_list.extend( - gathered[cp_rank][0, local_cu_seqlens[i] : local_cu_seqlens[i] + chunk_size] for cp_rank in range(cp_size) - ) - # Second half: rank cp_size-1 chunk, ..., rank 0 chunk (reversed) - whole_list.extend( - [ - gathered[cp_rank][0, local_cu_seqlens[i] + chunk_size : local_cu_seqlens[i + 1]] - for cp_rank in range(cp_size) - ][::-1] - ) - return torch.cat(whole_list).unsqueeze(0) # [1, T_global] - - -def _select_local_image_embeds( - full_input_ids: torch.Tensor, - cu_seqlens: torch.Tensor, - image_token_id: int, - image_embeds: torch.Tensor, - cp_rank: int, - cp_size: int, -) -> torch.Tensor: - """Select the subset of *image_embeds* that falls in this CP rank's chunk. - - With zigzag CP, each rank holds specific chunks of each packed sequence. - The vision encoder produces embeddings for ALL image tokens (ordered by - position in the full sequence). This function returns only the embeddings - whose positions land in the local chunk. - - Args: - full_input_ids: Reconstructed full input_ids [1, T_global]. - cu_seqlens: **Global** cumulative sequence lengths. - image_token_id: Token id for image placeholder tokens. - image_embeds: Vision embeddings [N_total, hidden] for all image tokens. - cp_rank: This rank's position in the CP group. - cp_size: Total number of CP ranks. - - Returns: - Subset of image_embeds [N_local, hidden] for this rank's image tokens. - """ - device = full_input_ids.device - full_flat = full_input_ids[0] # [T_global] - full_mask = full_flat == image_token_id - - # Build boolean mask over T_global marking this rank's positions - T_global = full_flat.shape[0] - rank_mask = torch.zeros(T_global, dtype=torch.bool, device=device) - - num_seqs = len(cu_seqlens) - 1 - for i in range(num_seqs): - seq_start = cu_seqlens[i].item() - seqlen = (cu_seqlens[i + 1] - cu_seqlens[i]).item() - chunk_size = seqlen // (2 * cp_size) - - # First-half chunk for this rank - first_start = seq_start + cp_rank * chunk_size - rank_mask[first_start : first_start + chunk_size] = True - - # Second-half chunk for this rank (reversed order) - second_end = seq_start + seqlen - cp_rank * chunk_size - rank_mask[second_end - chunk_size : second_end] = True - - # Image tokens that belong to this rank - local_image_mask = full_mask & rank_mask - n_local = local_image_mask.sum().item() - - if n_local == 0: - return image_embeds[:0] # empty slice preserving hidden dim - if n_local == image_embeds.shape[0]: - return image_embeds # all image tokens are on this rank - - # Map positions to indices in image_embeds via cumulative sum - image_cumsum = full_mask.long().cumsum(0) # 1-indexed - local_positions = local_image_mask.nonzero(as_tuple=True)[0] - embed_indices = image_cumsum[local_positions] - 1 - return image_embeds[embed_indices] - - -# --------------------------------------------------------------------------- -# Megatron VL Model -# --------------------------------------------------------------------------- -class Glm4vMoeVLModel(MegatronModule): - """GLM-4.6V vision-language model for Megatron training. - - Wraps an HF vision encoder (only on first PP stage) together with a - standard Megatron Core GPTModel configured for M-RoPE. - """ - - def __init__( - self, - language_transformer_config, - language_transformer_layer_spec, - hf_vision_config, - parallel_output: bool = True, - pre_process: bool = True, - post_process: bool = True, - ) -> None: - super().__init__(config=language_transformer_config) - - self.pre_process = pre_process - self.post_process = post_process - self.image_token_id = language_transformer_config.image_token_id - self.video_token_id = language_transformer_config.video_token_id - self.spatial_merge_size = language_transformer_config.spatial_merge_size - - self.share_embeddings_and_output_weights = False - - # Vision encoder -- only on the first pipeline stage - self.vision_model = None - if self.pre_process: - from transformers.models.glm4v_moe.modeling_glm4v_moe import Glm4vMoeVisionModel - - self.vision_model = Glm4vMoeVisionModel._from_config(hf_vision_config) - # Freeze vision encoder — not trained during RL - self.vision_model.requires_grad_(False) - self.vision_model.eval() - hook_hf_module_setattr_for_tp_grad_sync(self.vision_model) - if torch.cuda.is_available(): - self.vision_model = self.vision_model.to("cuda") - - # Language model -- standard Megatron GPT with M-RoPE - self.language_model = MCoreGPTModel( - config=language_transformer_config, - transformer_layer_spec=language_transformer_layer_spec, - vocab_size=language_transformer_config.vocab_size, - max_sequence_length=language_transformer_config.language_max_sequence_length, - parallel_output=parallel_output, - position_embedding_type="mrope", - rotary_percent=language_transformer_config.rotary_percent, - pre_process=self.pre_process, - post_process=self.post_process, - rotary_base=language_transformer_config.rotary_base, - fp16_lm_cross_entropy=language_transformer_config.fp16_lm_cross_entropy, - share_embeddings_and_output_weights=language_transformer_config.share_embeddings_and_output_weights, - scatter_embedding_sequence_parallel=False, - ) - - self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights - - # -- helpers required by Megatron pipeline engine ----------------------- - - def shared_embedding_or_output_weight(self): - return self.language_model.shared_embedding_or_output_weight() - - def set_input_tensor(self, input_tensor): - if not isinstance(input_tensor, list): - input_tensor = [input_tensor] - assert len(input_tensor) == 1 - if self.pre_process: - self.encoder_hidden_state = input_tensor[0] - else: - self.language_model.set_input_tensor(input_tensor[0]) - - # -- vision helpers ----------------------------------------------------- - - def _get_image_features(self, pixel_values, image_grid_thw): - """Run HF vision encoder and return flat image embeddings.""" - pixel_values = pixel_values.to(dtype=self.vision_model.dtype) - with torch.no_grad(): - return self.vision_model(pixel_values, grid_thw=image_grid_thw) - - # -- M-RoPE position IDs ----------------------------------------------- - - @staticmethod - def _get_vision_position_ids( - start_position: int, - grid_thw, - temp_merge_size: int, - spatial_merge_size: int, - device, - ) -> torch.Tensor: - """Compute 3D positions for one image/video (ported from HF).""" - llm_grid_t = grid_thw[0].item() // temp_merge_size - llm_grid_h = grid_thw[1].item() // spatial_merge_size - llm_grid_w = grid_thw[2].item() // spatial_merge_size - n_tokens = llm_grid_h * llm_grid_w * llm_grid_t - - pos_w = torch.arange(start_position, start_position + llm_grid_w, device=device) - pos_w = pos_w.repeat(llm_grid_h * llm_grid_t) - pos_h = torch.arange(start_position, start_position + llm_grid_h, device=device) - pos_h = pos_h.repeat_interleave(llm_grid_w * llm_grid_t) - pos_t = torch.full((n_tokens,), start_position, device=device, dtype=torch.long) - return torch.stack([pos_t, pos_h, pos_w], dim=0) # [3, n_tokens] - - def _compute_mrope_position_ids( - self, - input_ids_bshd: torch.Tensor, - image_grid_thw: torch.Tensor | None, - ) -> torch.Tensor: - """Compute 3D M-RoPE position IDs from input_ids in [bs, seq] format. - - Image regions are detected by looking for consecutive runs of - ``image_token_id`` in each sequence — no ``mm_token_type_ids`` needed. - """ - bs, seq_len = input_ids_bshd.shape - device = input_ids_bshd.device - spatial_merge_size = self.spatial_merge_size - - position_ids = torch.zeros(3, bs, seq_len, dtype=torch.long, device=device) - - if image_grid_thw is None or image_grid_thw.numel() == 0: - # Text-only: standard 1D positions replicated across 3 dims - pos = torch.arange(seq_len, device=device).unsqueeze(0).expand(bs, -1) - position_ids[0] = pos - position_ids[1] = pos - position_ids[2] = pos - return position_ids - - grid_iter = iter(image_grid_thw) - - for b in range(bs): - ids = input_ids_bshd[b] - is_image = ids == self.image_token_id - - # Find contiguous groups: text (0) vs image (1) - token_types = is_image.long() - groups = [] - for key, group in itertools.groupby(enumerate(token_types.tolist()), lambda x: x[1]): - g = list(group) - groups.append((key, g[0][0], g[-1][0] + 1)) - - current_pos = 0 - pos_list = [] - for modality, start, end in groups: - if modality == 0: - # Text tokens - n = end - start - pos_list.append(torch.arange(n, device=device).view(1, -1).expand(3, -1) + current_pos) - current_pos += n - else: - # Image tokens - grid_thw = next(grid_iter) - temp_merge_size = grid_thw[0] - vis_pos = self._get_vision_position_ids( - current_pos, - grid_thw, - temp_merge_size, - spatial_merge_size, - device, - ) - pos_list.append(vis_pos) - current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size - - all_pos = torch.cat(pos_list, dim=1) # [3, seq_for_this_sample] - position_ids[:, b, : all_pos.shape[1]] = all_pos - - return position_ids - - # -- forward ------------------------------------------------------------ - - def forward( - self, - input_ids: torch.Tensor, - position_ids: torch.Tensor = None, - attention_mask: torch.Tensor = None, - labels: torch.Tensor = None, - loss_mask: torch.Tensor = None, - inference_params=None, - packed_seq_params: PackedSeqParams = None, - extra_block_kwargs: dict = None, - # multimodal kwargs (unpacked from multimodal_train_inputs) - pixel_values: torch.Tensor = None, - image_grid_thw: torch.Tensor = None, - # unused VL kwargs that may come through - pixel_values_videos: torch.Tensor = None, - video_grid_thw: torch.Tensor = None, - mm_token_type_ids: torch.Tensor = None, - **kwargs, - ) -> torch.Tensor: - assert pixel_values_videos is None, "Video not supported yet" - assert inference_params is None, "Inference not supported" - - # -- Extract cu_seqlens and CP info early (needed for both vision scatter and M-RoPE) -- - cu_seqlens = None - if packed_seq_params is not None: - cu_seqlens = ( - packed_seq_params.cu_seqlens_q_padded - if packed_seq_params.cu_seqlens_q_padded is not None - else packed_seq_params.cu_seqlens_q - ) - cp_size = parallel_state.get_context_parallel_world_size() - full_input_ids = None # cached for reuse between vision scatter and M-RoPE - - combined_embeddings = None - - if self.pre_process: - # 1. Text embeddings from language model embedding layer - combined_embeddings = self.language_model.embedding( - input_ids=input_ids, - position_ids=None, - ).clone() # [seq, batch, hidden] - - # 2. Vision encoding + masked scatter - if pixel_values is not None and image_grid_thw is not None: - image_embeds = self._get_image_features(pixel_values, image_grid_thw) - image_embeds = image_embeds.to(combined_embeddings.device, combined_embeddings.dtype) - - # With CP > 1, input_ids is a local chunk but pixel_values - # cover ALL images. Select only the embeddings whose tokens - # land in this rank's zigzag portion. - if cp_size > 1 and cu_seqlens is not None: - full_input_ids = _gather_input_ids_from_cp(input_ids, cu_seqlens) - cp_rank = parallel_state.get_context_parallel_rank() - image_embeds = _select_local_image_embeds( - full_input_ids, - cu_seqlens, - self.image_token_id, - image_embeds, - cp_rank, - cp_size, - ) - - image_mask = (input_ids == self.image_token_id).contiguous() - # Scatter: [seq, bs, hidden] → [bs, seq, hidden] - combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() - if image_mask.any(): - combined_embeddings[image_mask] = image_embeds - combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() - - # Scatter to sequence-parallel region if needed - if self.config.sequence_parallel: - combined_embeddings = tensor_parallel.scatter_to_sequence_parallel_region(combined_embeddings) - combined_embeddings = combined_embeddings.contiguous() - - # 3. Compute M-RoPE position IDs - # position_ids must be available on ALL PP stages for rotary embeddings. - # On stage 0, compute from input_ids. Then broadcast to other stages. - pp_size = parallel_state.get_pipeline_model_parallel_world_size() - - if position_ids is None: - if self.pre_process: - # First PP stage: compute position_ids from input_ids. - # With CP > 1, input_ids is a local chunk; reconstruct full - # sequence so that _compute_mrope_position_ids sees all tokens - # (image token positions affect the M-RoPE IDs). - if cu_seqlens is not None: - if cp_size > 1: - if full_input_ids is None: - full_input_ids = _gather_input_ids_from_cp(input_ids, cu_seqlens) - else: - full_input_ids = input_ids - input_ids_bshd = _thd_to_bshd(full_input_ids, cu_seqlens) - pos_bshd = self._compute_mrope_position_ids(input_ids_bshd, image_grid_thw) - pos_packed = _bshd_to_thd(pos_bshd.permute(1, 2, 0), cu_seqlens) - position_ids = pos_packed.permute(2, 0, 1).contiguous() # [3, 1, T_global] - else: - position_ids = self._compute_mrope_position_ids(input_ids, image_grid_thw) - else: - # Non-first PP stage: allocate buffer with correct shape - if cu_seqlens is not None: - T = cu_seqlens[-1].item() - position_ids = torch.zeros(3, 1, T, dtype=torch.long, device=torch.cuda.current_device()) - else: - raise NotImplementedError( - "Non-THD position_ids broadcast not yet supported for non-first PP stages" - ) - - # Broadcast position_ids from first to all PP stages - if pp_size > 1: - src = parallel_state.get_pipeline_model_parallel_first_rank() - torch.distributed.broadcast( - position_ids, - src=src, - group=parallel_state.get_pipeline_model_parallel_group(), - ) - - # 4. Language model forward (pass decoder_input to skip re-embedding) - output = self.language_model( - input_ids=None, - position_ids=position_ids, - attention_mask=attention_mask, - decoder_input=combined_embeddings, - labels=labels, - loss_mask=loss_mask, - inference_params=inference_params, - packed_seq_params=packed_seq_params, - **(extra_block_kwargs or {}), - ) - - return output - - -# --------------------------------------------------------------------------- -# Model Provider (dataclass that doubles as TransformerConfig) -# --------------------------------------------------------------------------- -@dataclass -class Glm4vMoeVLModelProvider(GPTModelProvider): - """Provider that creates Glm4vMoeVLModel. - - Inherits from GPTModelProvider to reuse MoE + TransformerConfig infra. - Defined at module level (not inside a function) so that the class is - picklable -- megatron-bridge broadcasts config objects across PP ranks - via ``torch.distributed.broadcast_object_list`` which requires pickling. - """ - - # GLM-4.6V specific config - image_token_id: int = 151363 - video_token_id: int = 151364 - spatial_merge_size: int = 2 - - # Vision config (stored as HF config object) - hf_vision_config: object = None - hf_text_config: object = None - - # M-RoPE - position_embedding_type: str = "mrope" - mrope_section: list[int] = field(default_factory=lambda: [8, 12, 12]) - scatter_embedding_sequence_parallel: bool = False - - # Language model sequence length - language_max_sequence_length: int = 131072 - - def provide(self, pre_process=None, post_process=None, vp_stage=None): - """Create a Glm4vMoeVLModel instance.""" - - # Resolve PP stage flags - if pre_process is None: - pre_process = parallel_state.is_pipeline_first_stage(ignore_virtual=False, vp_stage=vp_stage) - if post_process is None: - post_process = parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage) - - # Build per-layer specs respecting moe_layer_freq (layer 0 = dense, rest = MoE) - transformer_layer_spec = get_gpt_decoder_block_spec( - config=self, - use_transformer_engine=True, - vp_stage=vp_stage, - ) - - model = Glm4vMoeVLModel( - language_transformer_config=self, - language_transformer_layer_spec=transformer_layer_spec, - hf_vision_config=self.hf_vision_config, - parallel_output=True, - pre_process=pre_process, - post_process=post_process, - ) - - return model - - -# --------------------------------------------------------------------------- -# Bridge -# --------------------------------------------------------------------------- -try: - from transformers import Glm4vMoeForConditionalGeneration as _Glm4vMoeHF -except ImportError: - _Glm4vMoeHF = "Glm4vMoeForConditionalGeneration" - - -@MegatronModelBridge.register_bridge(source=_Glm4vMoeHF, target=Glm4vMoeVLModel) -class Glm4vMoeBridge(MegatronModelBridge): - """Bridge between HuggingFace GLM-4.6V and the Megatron VL model.""" - - def provider_bridge(self, hf_pretrained): - """Create a Glm4vMoeVLModelProvider from HF config.""" - hf_config = hf_pretrained.config - text_config = hf_config.text_config - vision_config = deepcopy(hf_config.vision_config) - - model_dtype = self.dtype_from_hf(text_config, default=torch.bfloat16) - vision_config.torch_dtype = model_dtype - - ProviderClass = Glm4vMoeVLModelProvider - - rope_params = getattr(text_config, "rope_parameters", {}) or {} - mrope_section = rope_params.get("mrope_section", [8, 12, 12]) - rotary_base = rope_params.get("rope_theta", 500000) - partial_rotary_factor = rope_params.get("partial_rotary_factor", 0.5) - - # Determine MoE layer frequency - first_k_dense = getattr(text_config, "first_k_dense_replace", 1) - num_layers = text_config.num_hidden_layers - # Build moe_layer_freq list: first_k_dense dense layers + rest MoE - moe_layer_freq_list = [0] * first_k_dense + [1] * (num_layers - first_k_dense) - - # Shared expert intermediate size - n_shared = getattr(text_config, "n_shared_experts", 1) - moe_ffn = getattr(text_config, "moe_intermediate_size", 1408) - shared_expert_intermediate = moe_ffn * n_shared - - provider = ProviderClass( - # Language model configuration - num_layers=num_layers, - hidden_size=text_config.hidden_size, - ffn_hidden_size=text_config.intermediate_size, - num_attention_heads=text_config.num_attention_heads, - num_query_groups=text_config.num_key_value_heads, - kv_channels=getattr(text_config, "head_dim", 128), - init_method_std=text_config.initializer_range, - layernorm_epsilon=text_config.rms_norm_eps, - normalization="RMSNorm", - gated_linear_unit=True, - add_bias_linear=False, - hidden_dropout=0.0, - autocast_dtype=model_dtype, - make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(text_config.vocab_size), - rotary_base=rotary_base, - rotary_percent=partial_rotary_factor, - share_embeddings_and_output_weights=getattr(text_config, "tie_word_embeddings", False), - vocab_size=text_config.vocab_size, - seq_length=text_config.max_position_embeddings, - fp16=(model_dtype == torch.float16), - bf16=(model_dtype == torch.bfloat16), - params_dtype=model_dtype, - # MoE configuration - num_moe_experts=getattr(text_config, "n_routed_experts", 128), - moe_router_topk=getattr(text_config, "num_experts_per_tok", 8), - moe_ffn_hidden_size=moe_ffn, - moe_shared_expert_intermediate_size=shared_expert_intermediate, - moe_layer_freq=moe_layer_freq_list, - moe_grouped_gemm=True, - moe_token_dispatcher_type="alltoall", - moe_permute_fusion=True, - moe_router_load_balancing_type="seq_aux_loss", - moe_aux_loss_coeff=0, - moe_router_score_function="sigmoid", - moe_router_pre_softmax=True, - moe_router_enable_expert_bias=True, - moe_router_dtype="fp32", - # Attention - add_qkv_bias=getattr(text_config, "attention_bias", True), - qk_layernorm=getattr(text_config, "qk_layernorm", False) or getattr(text_config, "use_qk_norm", False), - # M-RoPE - mrope_section=mrope_section, - position_embedding_type="mrope", - scatter_embedding_sequence_parallel=False, - # Vision - hf_vision_config=vision_config, - hf_text_config=text_config, - image_token_id=getattr(hf_config, "image_token_id", 151363), - video_token_id=getattr(hf_config, "video_token_id", 151364), - spatial_merge_size=getattr(hf_config.vision_config, "spatial_merge_size", 2), - language_max_sequence_length=text_config.max_position_embeddings, - ) - - return provider - - def mapping_registry(self) -> MegatronMappingRegistry: - """Weight mappings from HF GLM-4.6V to Megatron format. - - Follows GLM-4.5 bridge pattern with language_model prefix for VL model. - Layer 0 is dense, layers 1-45 are MoE. The mapping framework handles - missing keys gracefully (warnings for non-existent params). - """ - param_mappings = { - # Embeddings and output - "language_model.embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight", - "language_model.output_layer.weight": "lm_head.weight", - "language_model.decoder.final_layernorm.weight": "model.language_model.norm.weight", - # Attention: input layernorm (fused with TE) - "language_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.language_model.layers.*.input_layernorm.weight", - # Attention: separate input layernorm (quantization layer spec) - "language_model.decoder.layers.*.input_layernorm.weight": "model.language_model.layers.*.input_layernorm.weight", - # Attention output - "language_model.decoder.layers.*.self_attention.linear_proj.weight": "model.language_model.layers.*.self_attn.o_proj.weight", - # Post-attention layernorm: - # MoE layers → pre_mlp_layernorm, Dense layer → linear_fc1.layer_norm_weight (fused) - "language_model.decoder.layers.*.pre_mlp_layernorm.weight": "model.language_model.layers.*.post_attention_layernorm.weight", - "language_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.language_model.layers.*.post_attention_layernorm.weight", - # Dense MLP output (layer 0) - "language_model.decoder.layers.*.mlp.linear_fc2.weight": "model.language_model.layers.*.mlp.down_proj.weight", - # MoE router - "language_model.decoder.layers.*.mlp.router.weight": "model.language_model.layers.*.mlp.gate.weight", - "language_model.decoder.layers.*.mlp.router.expert_bias": "model.language_model.layers.*.mlp.gate.e_score_correction_bias", - # MoE expert output (TEGroupedMLP format: weight* suffix) - "language_model.decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.language_model.layers.*.mlp.experts.*.down_proj.weight", - # MoE shared expert output - "language_model.decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.language_model.layers.*.mlp.shared_experts.down_proj.weight", - } - - mapping_list = [] - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - mapping_list.extend( - [ - # Vision model weights — replicated directly - ReplicatedMapping( - megatron_param="vision_model.**", - hf_param="model.visual.**", - ), - # QKV weight and bias - QKVMapping( - megatron_param="language_model.decoder.layers.*.self_attention.linear_qkv.weight", - q="model.language_model.layers.*.self_attn.q_proj.weight", - k="model.language_model.layers.*.self_attn.k_proj.weight", - v="model.language_model.layers.*.self_attn.v_proj.weight", - ), - QKVMapping( - megatron_param="language_model.decoder.layers.*.self_attention.linear_qkv.bias", - q="model.language_model.layers.*.self_attn.q_proj.bias", - k="model.language_model.layers.*.self_attn.k_proj.bias", - v="model.language_model.layers.*.self_attn.v_proj.bias", - ), - # Dense MLP gate+up (layer 0) - GatedMLPMapping( - megatron_param="language_model.decoder.layers.*.mlp.linear_fc1.weight", - gate="model.language_model.layers.*.mlp.gate_proj.weight", - up="model.language_model.layers.*.mlp.up_proj.weight", - ), - # MoE expert gate+up (TEGroupedMLP format) - GatedMLPMapping( - megatron_param="language_model.decoder.layers.*.mlp.experts.linear_fc1.weight*", - gate="model.language_model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.language_model.layers.*.mlp.experts.*.up_proj.weight", - ), - # MoE expert gate+up (SequentialMLP format, for quantization) - GatedMLPMapping( - megatron_param="language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", - gate="model.language_model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.language_model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", - hf_param="model.language_model.layers.*.mlp.experts.*.down_proj.weight", - ), - # MoE shared expert gate+up - GatedMLPMapping( - megatron_param="language_model.decoder.layers.*.mlp.shared_experts.linear_fc1.weight", - gate="model.language_model.layers.*.mlp.shared_experts.gate_proj.weight", - up="model.language_model.layers.*.mlp.shared_experts.up_proj.weight", - ), - ] - ) - - return MegatronMappingRegistry(*mapping_list) diff --git a/vime_plugins/models/flash_dot_product_attention.py b/vime_plugins/models/flash_dot_product_attention.py index b20feabb8..fdbc99afe 100644 --- a/vime_plugins/models/flash_dot_product_attention.py +++ b/vime_plugins/models/flash_dot_product_attention.py @@ -18,6 +18,7 @@ from megatron.core.utils import divide from torch import Tensor +from vime.utils import accelerator from vime_plugins.models.learnable_softmax_attention import learnable_softmax_flash_attn_varlen @@ -75,7 +76,7 @@ def __init__( elif config.softmax_type == "off-by-one": self.softmax_offset = torch.zeros( num_heads_per_partition, - device=torch.cuda.current_device(), + device=accelerator.current_device(), dtype=config.params_dtype, ) elif config.softmax_type == "learnable": @@ -84,7 +85,7 @@ def __init__( torch.nn.Parameter( torch.empty( num_heads_per_partition, - device=torch.cuda.current_device(), + device=accelerator.current_device(), dtype=config.params_dtype, ) ), diff --git a/vime_plugins/models/glm5/glm5.py b/vime_plugins/models/glm5/glm5.py index 273b5a745..5847ea6f6 100644 --- a/vime_plugins/models/glm5/glm5.py +++ b/vime_plugins/models/glm5/glm5.py @@ -1,5 +1,6 @@ import copy import math +import os from dataclasses import dataclass from typing import NoReturn @@ -27,7 +28,244 @@ from transformers import AutoConfig from .ops.indexer import generate_varlen_mask_params, lighting_indexer -from .ops.sparse_mla import SparseMLA +from .ops.sparse_mla import SparseMLA, VLLMSparseMLA + +# Names of the indexer submodules. On a DSA model with *cross-layer index +# sharing* these only exist on "computing" layers; "skip" layers drop them. +_INDEXER_SUBMODULE_NAMES = ("wq_b", "wk", "k_norm", "weights_proj") +_VLLM_ROPE_CACHE = {} + + +class _VLLMAbsorbWeightSTE(torch.autograd.Function): + @staticmethod + def forward(ctx, aligned_weight: torch.Tensor, trainable_weight: torch.Tensor): + if aligned_weight.shape != trainable_weight.shape: + raise ValueError( + "Aligned absorb-weight shape mismatch: " f"{aligned_weight.shape} != {trainable_weight.shape}" + ) + return aligned_weight + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + return None, grad_output + + +class _VLLMIndexerHeadWeights(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden_states: torch.Tensor, weight: torch.Tensor): + import deep_gemm + + flat_input = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + output = torch.empty( + (flat_input.shape[0], weight.shape[0]), + dtype=torch.float32, + device=flat_input.device, + ) + deep_gemm.bf16_gemm_nt( + flat_input, + weight.contiguous(), + output, + ) + ctx.save_for_backward(hidden_states, weight) + return output.view(*hidden_states.shape[:-1], weight.shape[0]) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + from vime.backends.megatron_utils.alignment.deepgemm_forward import router_gating_linear_backward + + hidden_states, weight = ctx.saved_tensors + return router_gating_linear_backward( + hidden_states, + weight, + grad_output, + torch.float32, + ) + + +def _get_vllm_indexer_head_weights( + hidden_states: torch.Tensor, + weight: torch.Tensor, + *, + num_heads: int, + head_dim: int, +) -> torch.Tensor: + if os.getenv("MEGATRON_USE_VLLM_FP8_INDEXER", "0") != "1": + output = WeightLinearFunction.apply(hidden_states, weight, None, torch.float32) + return output.squeeze(1) * (num_heads**-0.5) * (head_dim**-0.5) + + if hidden_states.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise TypeError( + "VLLM indexer head-gate alignment requires BF16 input and weight, " + f"got {hidden_states.dtype}/{weight.dtype}" + ) + output = _VLLMIndexerHeadWeights.apply(hidden_states, weight) + return output.squeeze(1) * (num_heads**-0.5) + + +def _get_fp8_aligned_absorb_weight(linear: torch.nn.Module) -> torch.Tensor: + weight = linear.weight + cache_key = ( + int(weight._version), + weight.data_ptr(), + weight.device, + weight.dtype, + tuple(weight.shape), + tuple(weight.stride()), + ) + aligned_weight = getattr(linear, "_vllm_fp8_aligned_absorb_weight_cache", None) + if aligned_weight is None or getattr(linear, "_vllm_fp8_aligned_absorb_weight_cache_key", None) != cache_key: + from vime.backends.megatron_utils.kernels.fp8_kernel import blockwise_cast_to_fp8_triton + + with torch.no_grad(): + from vllm.model_executor.layers.quantization.utils.fp8_utils import requant_weight_ue8m0_inplace + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used, per_block_cast_to_fp8 + + use_ue8m0 = is_deep_gemm_e8m0_used() + + if use_ue8m0: + # Blackwell (sm100/sm103): the VLLM rollout dequantizes the + # absorb weight from a UE8M0 (power-of-two scale) FP8 tensor. The + # Hopper blockwise_cast_to_fp8_triton (FP32 scales) produces a + # bf16 weight that differs by ~1 bf16 ULP from the UE8M0 dequant, + # which propagates as a uniform ~1-ULP offset in the absorbed q + # and, amplified through the layers and the FP32 LM head, shows + # up as a large train/rollout logprob gap. Replicate the UE8M0 + # quant so the dequantized absorb weight bit-matches the rollout. + # requant_weight_ue8m0_inplace updates both tensors to the raw + # power-of-two per-block scale used for this dequantization. + qweight, scale_inv = per_block_cast_to_fp8(weight.detach().contiguous(), block_size=(128, 128)) + requant_weight_ue8m0_inplace(qweight, scale_inv, (128, 128)) + else: + # Hopper (sm90): FP32 block scales. Unchanged. + qweight, scale_inv = blockwise_cast_to_fp8_triton( + weight.detach().contiguous(), + (128, 128), + ) + expanded_scale = scale_inv.repeat_interleave(128, dim=-2).repeat_interleave(128, dim=-1) + aligned_weight = (qweight.float() * expanded_scale[: qweight.shape[-2], : qweight.shape[-1]]).to( + torch.bfloat16 + ) + linear._vllm_fp8_aligned_absorb_weight_cache = aligned_weight + linear._vllm_fp8_aligned_absorb_weight_cache_key = cache_key + return _VLLMAbsorbWeightSTE.apply(aligned_weight, weight) + + +@torch.no_grad() +def _apply_vllm_rope_forward( + value: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + from vllm import _custom_ops as ops + + output = torch.empty_strided(value.size(), value.stride(), dtype=value.dtype, device=value.device) + output.copy_(value) + ops.rotary_embedding( + positions, + output, + None, + value.shape[-1], + cos_sin_cache, + False, + ) + return output + + +class _VLLMRoPE(torch.autograd.Function): + @staticmethod + def forward(ctx, value, cos_sin_cache, positions): + ctx.save_for_backward(cos_sin_cache, positions) + return _apply_vllm_rope_forward(value, cos_sin_cache, positions) + + @staticmethod + def backward(ctx, grad_output): + cos_sin_cache, positions = ctx.saved_tensors + half = grad_output.shape[-1] // 2 + broadcast_shape = (positions.numel(),) + (1,) * (grad_output.ndim - 2) + (half,) + cos = cos_sin_cache[positions, :half].view(broadcast_shape).to(grad_output.dtype) + sin = cos_sin_cache[positions, half:].view(broadcast_shape).to(grad_output.dtype) + grad_even = grad_output[..., 0::2] + grad_odd = grad_output[..., 1::2] + grad_input = torch.stack( + ( + grad_even * cos + grad_odd * sin, + grad_odd * cos - grad_even * sin, + ), + dim=-1, + ).flatten(-2) + return grad_input, None, None + + +def _get_vllm_rope_cache( + device: torch.device, + rotary_dim: int, + rotary_base: float, + needed_positions: int, +) -> torch.Tensor: + key = (device.type, device.index, rotary_dim, float(rotary_base)) + cache = _VLLM_ROPE_CACHE.get(key) + if cache is not None and cache.shape[0] >= needed_positions: + return cache + inv_freq = 1.0 / (rotary_base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) / rotary_dim)) + positions = torch.arange(needed_positions, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j -> ij", positions, inv_freq) + cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1) + _VLLM_ROPE_CACHE[key] = cache + return cache + + +class _DSAKVFP8QAT(torch.autograd.Function): + @staticmethod + def forward(ctx, kv: torch.Tensor): + if kv.dtype != torch.bfloat16 or kv.shape[-2:] != (1, 576): + raise ValueError( + "GLM5 KV FP8 QAT requires BF16 [..., 1, 576], " f"got dtype={kv.dtype}, shape={tuple(kv.shape)}" + ) + flat_kv = kv.contiguous().view(-1, 576) + nope = flat_kv[:, :512].view(-1, 4, 128).float() + scale = nope.abs().amax(dim=-1, keepdim=True).clamp_min(1e-10) / torch.finfo(torch.float8_e4m3fn).max + quantized = ( + (nope / scale) + .clamp( + -torch.finfo(torch.float8_e4m3fn).max, + torch.finfo(torch.float8_e4m3fn).max, + ) + .to(torch.float8_e4m3fn) + ) + dequantized_nope = (quantized.float() * scale).to(torch.bfloat16).view(-1, 512) + return torch.cat((dequantized_nope, flat_kv[:, 512:]), dim=-1).view_as(kv) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + return grad_output + + +def _fake_quant_fp8_kv_cache(kv: torch.Tensor) -> torch.Tensor: + enabled = os.environ.get("DSA_KV_FP8_QAT", "0").strip().lower() + if enabled not in ("1", "true", "yes", "on"): + return kv + block_size = int(os.environ.get("DSA_KV_FP8_QAT_BLOCK_SIZE", "128")) + if block_size != 128: + raise ValueError("VLLM-aligned KV FP8 QAT requires block size 128") + return _DSAKVFP8QAT.apply(kv) + + +def is_skip_topk_layer(layer_number: int, skip_topk_offset: int, topk_freq: int) -> bool: + """Whether the (1-indexed) Megatron ``layer_number`` reuses a previous layer's top-k. + + Mirrors ``glm-train-prod``'s ``_get_skip_topk_flags``: a layer *computes* its + own top-k when ``max(layer_number - offset, 0) % freq == 0``; otherwise it is a + skip layer that reuses the most recent computing layer's indices. + """ + return (max(layer_number - skip_topk_offset, 0) % topk_freq) != 0 + + +def source_compute_layer(layer_number: int, skip_topk_offset: int, topk_freq: int) -> int: + """The computing layer whose ``topk_indices`` a skip layer reuses.""" + layer = layer_number + while is_skip_topk_layer(layer, skip_topk_offset, topk_freq): + layer -= 1 + return layer @dataclass @@ -137,6 +375,30 @@ def __init__( self.index_topk = 2048 + # Cross-layer index sharing (optional). When the HF config provides + # ``index_topk_freq`` / ``index_skip_topk_offset`` (see ``get_glm5_spec``), + # only a subset of "computing" layers run the indexer top-k; the remaining + # "skip" layers reuse the most recent computing layer's ``topk_indices``. + # When those attrs are absent (``freq`` defaults to 1) every layer computes + # its own top-k and ``skip_topk`` is always False -- i.e. the plain DSA path. + self.index_topk_freq = getattr(config, "index_topk_freq", 1) or 1 + self.skip_topk_offset = getattr(config, "index_skip_topk_offset", 0) or 0 + self.index_share = self.index_topk_freq > 1 + self.skip_topk = self.index_share and is_skip_topk_layer( + layer_number, self.skip_topk_offset, self.index_topk_freq + ) + self._source_layer = ( + source_compute_layer(layer_number, self.skip_topk_offset, self.index_topk_freq) + if self.index_share + else layer_number + ) + + # Attribute name of the per-microbatch top-k holder we attach to the + # ``packed_seq_params`` object (a plain dict: source layer_number -> topk_indices). + # Used only on index-share models; see ``forward`` for why it lives on + # ``packed_seq_params`` (per-microbatch isolation + recompute safety under PP). + _HOLDER_ATTR = "_dsa_index_share_topk_holder" + def forward( self, hidden_states, @@ -204,15 +466,62 @@ def fused_select_topk(index_q, index_k, w, starts, ends, block_size=8192): topk_indices.append(topk_indices_block) return torch.cat(indexer_topk_scores, dim=0), torch.cat(topk_indices, dim=0).unsqueeze(1) - starts, ends = generate_varlen_mask_params(packed_seq_params.cu_seqlens_q) - index_key = index_key.squeeze(1) - head_weights = head_weights.unsqueeze(-1) - - starts = scatter_to_sequence_parallel_region(starts, group=parallel_state.get_context_parallel_group()) - ends = scatter_to_sequence_parallel_region(ends, group=parallel_state.get_context_parallel_group()) - _, topk_indices = fused_select_topk(index_query, index_key, head_weights, starts, ends) - - core_attn_out, _ = SparseMLA.apply(q, kv, topk_indices, self.softmax_scale) + if self.index_share: + # Cross-layer index sharing. The top-k holder lives on the per-microbatch + # ``packed_seq_params`` object: it is constructed fresh per microbatch in + # ``get_batch`` and is closure-captured by Megatron's activation-checkpoint + # ``custom_forward``, so the same instance is reused at recompute time. + # That gives per-microbatch isolation (no cross-microbatch clobber under + # PP 1F1B) AND recompute safety (the computing layer's entry written in the + # original forward is still present when a skip layer's chunk recomputes). + # Note: this never crosses a PP boundary -- a stage always starts on a + # computing layer (asserted in ``get_glm5_spec``), so a skip layer's source + # is always in-stage. + holder = getattr(packed_seq_params, self._HOLDER_ATTR, None) + if holder is None: + holder = {} + setattr(packed_seq_params, self._HOLDER_ATTR, holder) + + if self.skip_topk: + if self._source_layer not in holder: + raise AssertionError( + "DSA index-share: skip layer " + f"(layer_number={self.layer_number}) needs the top-k of its source " + f"computing layer (layer_number={self._source_layer}), but that layer " + "did not run in this pipeline stage's forward. Cross-PP top-k sharing " + "is not supported; ensure every pipeline stage starts on a computing " + f"layer (index_topk_freq={self.index_topk_freq}, " + f"index_skip_topk_offset={self.skip_topk_offset}). " + f"Holder has layers {sorted(holder)}." + ) + topk_indices = holder[self._source_layer] + else: + starts, ends = generate_varlen_mask_params(packed_seq_params.cu_seqlens_q) + index_key = index_key.squeeze(1) + head_weights = head_weights.unsqueeze(-1) + starts = scatter_to_sequence_parallel_region(starts, group=parallel_state.get_context_parallel_group()) + ends = scatter_to_sequence_parallel_region(ends, group=parallel_state.get_context_parallel_group()) + _, topk_indices = fused_select_topk(index_query, index_key, head_weights, starts, ends) + holder[self.layer_number] = topk_indices + else: + starts, ends = generate_varlen_mask_params(packed_seq_params.cu_seqlens_q) + index_key = index_key.squeeze(1) + head_weights = head_weights.unsqueeze(-1) + + starts = scatter_to_sequence_parallel_region(starts, group=parallel_state.get_context_parallel_group()) + ends = scatter_to_sequence_parallel_region(ends, group=parallel_state.get_context_parallel_group()) + _, topk_indices = fused_select_topk(index_query, index_key, head_weights, starts, ends) + + if os.getenv("MEGATRON_USE_VLLM_SPARSE_MLA", "0") == "1": + core_attn_out, _ = VLLMSparseMLA.apply( + q, + kv, + topk_indices, + self.softmax_scale, + self.config.kv_lora_rank, + ) + else: + core_attn_out, _ = SparseMLA.apply(q, kv, topk_indices, self.softmax_scale) core_attn_out = torch.einsum("thm,hdm->thd", core_attn_out, wv) core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) @@ -246,6 +555,7 @@ def __init__( cp_comm_type: str | None = None, model_comm_pgs=None, pg_collection=None, + name: str | None = None, ): super().__init__( config=config, @@ -403,6 +713,50 @@ def __init__( ) self.weights_proj.weight._skip_gather = True + # Index-share skip layers carry no indexer weights -- drop the modules the + # base path built unconditionally so the parameter set matches the + # checkpoint (which only stores indexer weights on computing layers) and so + # weight export to HF naturally omits them on skip layers. + if self.skip_topk: + for name in _INDEXER_SUBMODULE_NAMES: + if hasattr(self, name): + delattr(self, name) + + @torch.no_grad() + def _get_indexer_q_input(self, q_compressed: torch.Tensor) -> torch.Tensor: + """Return the q-RMSNorm value consumed by the DSA indexer. + + GLM-5 fuses q RMSNorm into ``linear_q_up_proj``. That fused projection + does not update its input, so the indexer must reconstruct the same + normalized q-latent before applying ``wq_b``. + """ + q_compressed = q_compressed.detach() + fused_norm_weight = getattr(self.linear_q_up_proj, "layer_norm_weight", None) + if fused_norm_weight is None: + if isinstance(self.q_layernorm, IdentityOp): + return q_compressed + return self.q_layernorm(q_compressed) + + if self.config.normalization != "RMSNorm": + raise ValueError(f"GLM-5 DSA indexer expects RMSNorm, got {self.config.normalization}") + norm_weight = fused_norm_weight.detach().float() + if self.config.layernorm_zero_centered_gamma: + norm_weight = norm_weight + 1.0 + if os.getenv("MEGATRON_USE_VLLM_FUSED_RESIDUAL_RMS", "0") == "1": + from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant + + return rms_norm_batch_invariant( + q_compressed, + norm_weight, + self.config.layernorm_epsilon, + ) + return torch.nn.functional.rms_norm( + q_compressed.float(), + normalized_shape=(q_compressed.shape[-1],), + weight=norm_weight, + eps=self.config.layernorm_epsilon, + ).to(q_compressed.dtype) + def get_absorb_query_key_value_tensors( self, hidden_states, @@ -427,8 +781,11 @@ def get_absorb_query_key_value_tensors( rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( inference_context, None, hidden_states, self.config, packed_seq_params ) - # TODO: support apply_rope_fusion - rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) + # YarnRotaryEmbedding/RotaryEmbedding.forward is wrapped in lru_cache, so it + # only accepts hashable args: pass the packed-sequence flag, not the + # (unhashable) PackedSeqParams object. + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == "thd" + rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) cu_seqlens_q = packed_seq_params.cu_seqlens_q cu_seqlens_kv = packed_seq_params.cu_seqlens_kv @@ -440,6 +797,7 @@ def get_absorb_query_key_value_tensors( # down proj are `TELinear`s, so the output is gathered and not TP-partitioned.` q_compressed, _ = self.linear_q_down_proj(hidden_states) q_compressed = q_compressed.squeeze(1) + index_q_input = None if self.skip_topk else self._get_indexer_q_input(q_compressed) kv_combined, _ = self.linear_kv_down_proj(hidden_states) if self.config.sequence_parallel: @@ -457,7 +815,12 @@ def get_absorb_query_key_value_tensors( q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) q_no_pe, q_pos_emb = torch.split(q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1) - w_kc, w_vc = self.linear_kv_up_proj.weight.unflatten( + absorb_weight = ( + _get_fp8_aligned_absorb_weight(self.linear_kv_up_proj) + if os.getenv("MEGATRON_USE_VLLM_SPARSE_MLA", "0") == "1" + else self.linear_kv_up_proj.weight + ) + w_kc, w_vc = absorb_weight.unflatten( 0, (-1, self.config.qk_head_dim + self.config.v_head_dim), ).split([self.config.qk_head_dim, self.config.v_head_dim], dim=1) @@ -479,6 +842,24 @@ def get_absorb_query_key_value_tensors( ) def fuse_rope(q, cu_seqlens, gathered=False): + if os.getenv("MEGATRON_USE_VLLM_ROPE", "0") == "1": + if parallel_state.get_tensor_model_parallel_world_size() != 1: + raise RuntimeError("MEGATRON_USE_VLLM_ROPE requires TP=1") + if parallel_state.get_context_parallel_world_size() != 1: + raise RuntimeError("The Vime GLM5 vLLM RoPE alignment currently requires CP=1") + token_ids = torch.arange(q.shape[0], dtype=torch.int64, device=q.device) + seq_ids = torch.searchsorted(cu_seqlens[1:], token_ids, right=True) + positions = token_ids - cu_seqlens[seq_ids] + cache = _get_vllm_rope_cache( + q.device, + q.shape[-1], + self.config.rotary_base, + int(positions.max().item()) + 1, + ) + if torch.is_grad_enabled() and q.requires_grad: + return _VLLMRoPE.apply(q, cache, positions) + return _apply_vllm_rope_forward(q, cache, positions) + # worse precision than apex. # from megatron.core.extensions.transformer_engine import fused_apply_rotary_pos_emb_thd from apex.transformer.functional import fused_apply_rotary_pos_emb_thd @@ -504,19 +885,24 @@ def fuse_rope(q, cu_seqlens, gathered=False): query = torch.cat([q_no_pe, q_pos_emb], dim=-1) key = torch.cat([kv_compressed, k_pos_emb], dim=-1) + key = _fake_quant_fp8_kv_cache(key) query = query.contiguous() key = key.contiguous() + if self.skip_topk: + # Index-share skip layer: reuse a previous layer's top-k, so the indexer + # projections are not run here. Return None for the index tensors. + return query, key, w_vc, None, None, None + # ========================================= # Indexer # ========================================= # Project queries and keys - q_compressed = q_compressed.detach() hidden_states = hidden_states.detach() rotary_pos_emb = rotary_pos_emb.detach() - index_q, _ = self.wq_b(q_compressed) + index_q, _ = self.wq_b(index_q_input) index_q = index_q.view( *index_q.size()[:-1], self.config.index_num_attention_heads, self.config.index_head_dim ) # [total_tokens, index_num_attention_heads_per_partition, index_head_dim] @@ -531,11 +917,12 @@ def fuse_rope(q, cu_seqlens, gathered=False): index_k = gather_from_sequence_parallel_region(index_k, group=parallel_state.get_context_parallel_group()) index_k = index_k.unsqueeze(1) # [total_tokens, 1, head_dim] - # head_weights, _ = self.weights_proj(hidden_states.float()) - head_weights = WeightLinearFunction.apply(hidden_states, self.weights_proj.weight, None, torch.float32) - head_weights = head_weights.squeeze(1) * ( - (self.config.index_num_attention_heads**-0.5) * (self.config.index_head_dim**-0.5) - ) # [total_tokens, index_num_attention_heads_per_partition] + head_weights = _get_vllm_indexer_head_weights( + hidden_states, + self.weights_proj.weight, + num_heads=self.config.index_num_attention_heads, + head_dim=self.config.index_head_dim, + ) if self.config.sequence_parallel: head_weights = gather_from_sequence_parallel_region(head_weights) @@ -607,6 +994,13 @@ def get_glm5_spec(args, config, vp_stage): hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) config.index_num_attention_heads = hf_config.index_n_heads config.index_head_dim = hf_config.index_head_dim + # Optional cross-layer index-sharing schedule. Present on DSA checkpoints that + # only store indexer weights on a subset of "computing" layers (e.g. GLM-5.2 + # 744B-A40B). When absent, every layer computes its own top-k (plain + # DSA) and ``DSAMLASelfAttention`` runs the non-shared path. + config.index_topk_freq = getattr(hf_config, "index_topk_freq", 1) or 1 + config.index_skip_topk_offset = getattr(hf_config, "index_skip_topk_offset", 0) or 0 + # Define the decoder block spec kwargs = { "use_transformer_engine": True, @@ -615,6 +1009,32 @@ def get_glm5_spec(args, config, vp_stage): kwargs["vp_stage"] = vp_stage transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs) num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) + + # Cross-layer index sharing keeps the shared top-k in a per-microbatch holder + # on ``packed_seq_params``, which does not cross PP boundaries. A skip layer + # therefore must run in the same pipeline stage as the computing layer it + # reuses. If a (virtual) pipeline stage *starts* with a skip layer, its source + # computing layer lives on a previous stage and the lookup would miss. Forbid + # that split here (supporting it would need PP send/recv of the top-k). + if config.index_topk_freq > 1: + from megatron.core.transformer.transformer_block import get_transformer_layer_offset + + layer_offset = get_transformer_layer_offset(config, vp_stage=vp_stage) + for local_id in range(num_layers_to_build): + layer_number = local_id + layer_offset + 1 # Megatron layer_number is 1-indexed + if local_id == 0 and is_skip_topk_layer( + layer_number, config.index_skip_topk_offset, config.index_topk_freq + ): + src = source_compute_layer(layer_number, config.index_skip_topk_offset, config.index_topk_freq) + raise AssertionError( + "DSA index-share pipeline split is invalid: this stage starts at global " + f"layer_number={layer_number} which is a skip layer whose source computing " + f"layer={src} is on a previous pipeline stage. Cross-layer top-k sharing does " + "not cross PP boundaries. Choose a pipeline layout where every stage begins on " + "a computing layer (index_topk_freq=" + f"{config.index_topk_freq}, index_skip_topk_offset={config.index_skip_topk_offset})." + ) + backend = TESpecProvider() self_attn_module_spec = ModuleSpec( diff --git a/vime_plugins/models/glm5/ops/indexer.py b/vime_plugins/models/glm5/ops/indexer.py index 210196b9b..a63c992ff 100644 --- a/vime_plugins/models/glm5/ops/indexer.py +++ b/vime_plugins/models/glm5/ops/indexer.py @@ -1,9 +1,96 @@ +import os + import torch +# Bind flashinfer.comm's module-level `cudart = CudaRTLibrary()` to the *real* +# libcudart before tilelang loads its `libcudart_stub.so`. On a cu12 box whose +# flashinfer pulled in cu13 deps, importing flashinfer.comm *after* the stub is +# resident makes `find_loaded_library("libcudart")` match the stub (which lacks +# `cudaDeviceReset`), crashing the lazy `dsa_indexer` import inside the forward. +# Importing it here — while only the real libcudart is loaded — caches the +# module with a valid binding; later imports are no-ops. Best-effort: other +# environments without flashinfer must still import this module. +try: # noqa: SIM105 + import flashinfer.comm # noqa: F401 +except Exception: + pass + from .tilelang_indexer_bwd import indexer_bwd_interface from .tilelang_indexer_fwd import indexer_fwd_interface +def _vllm_fp8_indexer_logits( + index_q: torch.Tensor, + index_k: torch.Tensor, + weights: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, +) -> torch.Tensor: + """Evaluate GLM5 indexer logits with VLLM's FP8 DSA kernels.""" + + if index_q.shape[-1] != 128: + raise ValueError("VLLM FP8 indexer alignment requires head_dim=128, " f"got {index_q.shape[-1]}") + from vllm import _custom_ops as ops + from vllm.model_executor.layers.quantization.utils.fp8_utils import per_token_group_quant_fp8 + from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + + q_rotated = torch.cat((index_q[..., -64:], index_q[..., :-64]), dim=-1).contiguous() + if index_k.ndim == 3: + if index_k.shape[1] != 1: + raise ValueError(f"Expected one indexer KV head, got {index_k.shape}") + index_k = index_k.squeeze(1) + k_rotated = torch.cat((index_k[..., -64:], index_k[..., :-64]), dim=-1).contiguous() + + q_fp8, q_scale = per_token_group_quant_fp8( + q_rotated.view(-1, q_rotated.shape[-1]), + 128, + use_ue8m0=True, + ) + q_fp8 = q_fp8.view_as(q_rotated) + q_scale = q_scale.view(*q_rotated.shape[:-1], 1) + page_size = 64 + num_k = k_rotated.shape[0] + num_pages = (num_k + page_size - 1) // page_size + packed_k = torch.empty( + (num_pages, page_size, 128 + 4), + dtype=torch.uint8, + device=k_rotated.device, + ) + ops.indexer_k_quant_and_cache( + k_rotated, + packed_k, + torch.arange(num_k, dtype=torch.int64, device=k_rotated.device), + 128, + "ue8m0", + ) + k_bytes = torch.empty((num_k, 128), dtype=torch.uint8, device=k_rotated.device) + k_scale_bytes = torch.empty((num_k, 4), dtype=torch.uint8, device=k_rotated.device) + ops.cp_gather_indexer_k_quant_cache( + packed_k, + k_bytes, + k_scale_bytes, + torch.arange(num_pages, dtype=torch.int32, device=k_rotated.device).unsqueeze(0), + torch.tensor([0, num_k], dtype=torch.int32, device=k_rotated.device), + ) + k_fp8 = k_bytes.view(torch.float8_e4m3fn) + k_scale = k_scale_bytes.view(torch.float32).reshape(-1) + scaled_weights = weights.float() * q_scale.squeeze(-1).float() + scaled_weights = (scaled_weights * (index_q.shape[-1] ** -0.5)).contiguous() + logits = fp8_fp4_mqa_logits( + (q_fp8, None), + (k_fp8, k_scale), + scaled_weights, + starts.to(torch.int32).contiguous(), + ends.to(torch.int32).contiguous(), + clean_logits=False, + ) + key_positions = torch.arange(num_k, dtype=torch.int32, device=index_q.device).unsqueeze(0) + valid = (key_positions >= starts.to(torch.int32).unsqueeze(1)) & ( + key_positions < ends.to(torch.int32).unsqueeze(1) + ) + return logits.masked_fill(~valid, float("-inf")) + + def pytorch_extract_topk_scores(logits, topk_indices, dim=-1): valid_mask = topk_indices != -1 safe_indices = topk_indices.clamp(min=0).to(torch.int64) @@ -12,6 +99,41 @@ def pytorch_extract_topk_scores(logits, topk_indices, dim=-1): return scores +def pytorch_topk_with_invalid_padding(logits: torch.Tensor, topk: int): + """Select up to ``topk`` keys and retain the fixed-width DSA layout. + + Short packed microbatches can contain fewer total keys than the model's + configured DSA top-k. VLLM represents the missing selections with -1; + mirror that representation instead of passing an out-of-range ``k`` to + ``torch.topk``. + """ + selected = min(topk, logits.shape[-1]) + scores, indices = torch.topk(logits, selected, dim=-1) + indices = indices.to(torch.int32) + indices = indices.masked_fill(scores == -torch.inf, -1) + if selected == topk: + return scores, indices + + pad_shape = (*logits.shape[:-1], topk - selected) + scores = torch.cat( + (scores, logits.new_full(pad_shape, float("-inf"))), + dim=-1, + ) + indices = torch.cat( + ( + indices, + torch.full( + pad_shape, + -1, + dtype=torch.int32, + device=logits.device, + ), + ), + dim=-1, + ) + return scores, indices + + class IndexerFunction(torch.autograd.Function): @staticmethod def forward( @@ -25,11 +147,32 @@ def forward( topk_indices: torch.Tensor | None = None, ): _, head_num, _ = index_q.shape - logits = indexer_fwd_interface(index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=True) + if os.getenv("MEGATRON_USE_VLLM_FP8_INDEXER", "0") == "1": + logits = _vllm_fp8_indexer_logits( + index_q, + index_k, + weights, + cu_seqlen_ks, + cu_seqlen_ke, + ) + else: + logits = indexer_fwd_interface( + index_q, + index_k, + weights, + cu_seqlen_ks, + cu_seqlen_ke, + clean_logits=True, + ) if topk_indices is None: - index_score, topk_indices = torch.topk(logits, topk, dim=-1) - topk_indices = topk_indices.to(torch.int32) - topk_indices = topk_indices.masked_fill(index_score == -torch.inf, -1) + index_score, topk_indices = pytorch_topk_with_invalid_padding(logits, topk) + if os.getenv("MEGATRON_USE_VLLM_FP8_INDEXER", "0") == "1": + invalid_sort_key = torch.iinfo(torch.int32).max + topk_indices = torch.sort( + topk_indices.masked_fill(topk_indices < 0, invalid_sort_key), + dim=-1, + ).values + topk_indices = topk_indices.masked_fill(topk_indices == invalid_sort_key, -1) index_score = pytorch_extract_topk_scores(logits, topk_indices) diff --git a/vime_plugins/models/glm5/ops/sparse_mla.py b/vime_plugins/models/glm5/ops/sparse_mla.py index 260b38ed0..ada1e8fbf 100644 --- a/vime_plugins/models/glm5/ops/sparse_mla.py +++ b/vime_plugins/models/glm5/ops/sparse_mla.py @@ -42,3 +42,66 @@ def backward(ctx, grad_output, grad_lse): # Return gradients for each input (None for indices as it's not differentiable) return tl_dq, tl_dkv, None, None + + +class VLLMSparseMLA(torch.autograd.Function): + """VLLM FlashMLA forward with the trainable TileLang backward.""" + + @staticmethod + def forward(ctx, q, kv, indices, scaling, d_v=512): + from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd + + q = q.contiguous() + kv = kv.contiguous() + indices = indices.contiguous() + + # flash_mla_sparse requires num_heads to be a multiple of 64 on Hopper + # (sm90) and 128 on Blackwell (sm100/sm103). The kernel is NOT + # padding-invariant on sm103: padding q from 64 -> 128 heads changes the + # bf16 rounding of the real heads (~1 bf16 ULP). The VLLM rollout + # (dsa_backend._forward_flashmla_sparse) always applies this padding on + # Blackwell, so the train side MUST pad identically or train/rollout + # logprobs diverge (0.027 on B300 vs 1.9e-7 on H100). Hopper needs no + # padding for 64 heads (64 % 64 == 0), so this branch is a no-op there. + num_heads = q.shape[1] + required_padding = 128 if torch.cuda.get_device_capability(q.device)[0] >= 10 else 64 + need_padding = num_heads % required_padding != 0 + if need_padding: + assert required_padding % num_heads == 0, ( + f"flash_mla_sparse num_heads {num_heads} cannot be padded to " f"{required_padding}" + ) + q_input = q.new_zeros((q.shape[0], required_padding, q.shape[2])) + q_input[:, :num_heads, :] = q + else: + q_input = q + + output, _, lse = flash_mla_sparse_fwd( + q=q_input, + kv=kv, + indices=indices, + sm_scale=scaling, + d_v=d_v, + ) + if need_padding: + output = output[:, :num_heads, :].contiguous() + lse = lse[:, :num_heads].contiguous() + ctx.scaling = scaling + ctx.d_v = d_v + ctx.save_for_backward(q, kv, indices, output, lse.contiguous()) + return output, lse.to(torch.bfloat16) + + @staticmethod + def backward(ctx, grad_output, grad_lse): + q, kv, indices, output, lse = ctx.saved_tensors + if grad_output is None: + grad_output = torch.zeros_like(output) + dq, dkv = sparse_mla_bwd( + q, + kv, + output, + grad_output.contiguous(), + indices, + lse, + sm_scale=ctx.scaling, + ) + return dq, dkv, None, None, None diff --git a/vime_plugins/models/gpt_oss.py b/vime_plugins/models/gpt_oss.py deleted file mode 100644 index 52ca4f1e0..000000000 --- a/vime_plugins/models/gpt_oss.py +++ /dev/null @@ -1,53 +0,0 @@ -"""GPT-OSS 20B model spec for Megatron. - -Replaces core_attention with FlashDotProductAttention to support -learnable softmax (attention sinks) + sliding window attention in -packed sequence (THD) format, which TE does not support. - -Also registers FlashDotProductAttention with megatron-bridge's AutoMapping -so the weight converter knows its parallelism type. - -Usage: - --spec "vime_plugins.models.gpt_oss" "get_gpt_oss_spec" -""" - -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec - -from vime_plugins.models.flash_dot_product_attention import FlashDotProductAttention - - -def _replace_core_attention_in_spec(spec, replacement_cls): - """Recursively replace core_attention in a layer/block spec.""" - if hasattr(spec, "layer_specs") and not hasattr(spec, "submodules"): - for layer_spec in spec.layer_specs: - _replace_core_attention_in_spec(layer_spec, replacement_cls) - return - if hasattr(spec, "submodules"): - sub = spec.submodules - if hasattr(sub, "core_attention"): - sub.core_attention = replacement_cls - if hasattr(sub, "layer_specs"): - for layer_spec in sub.layer_specs: - _replace_core_attention_in_spec(layer_spec, replacement_cls) - for attr in dir(sub): - if attr.startswith("_") or attr == "layer_specs": - continue - val = getattr(sub, attr) - if hasattr(val, "submodules"): - _replace_core_attention_in_spec(val, replacement_cls) - - -def get_gpt_oss_spec(args, config, vp_stage): - kwargs = {"use_transformer_engine": True} - if vp_stage is not None: - kwargs["vp_stage"] = vp_stage - transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs) - - _replace_core_attention_in_spec(transformer_layer_spec, FlashDotProductAttention) - - # Register with megatron-bridge so weight converter knows the parallelism type. - from megatron.bridge.models.conversion.param_mapping import AutoMapping - - AutoMapping.register_module_type("FlashDotProductAttention", "column") - - return transformer_layer_spec diff --git a/vime_plugins/models/qwen3_5.py b/vime_plugins/models/qwen3_5.py index 294c9d97a..01c81dc8c 100644 --- a/vime_plugins/models/qwen3_5.py +++ b/vime_plugins/models/qwen3_5.py @@ -9,6 +9,8 @@ from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from transformers.activations import ACT2FN +from vime.utils import accelerator + try: from fla.modules import FusedRMSNormGated, ShortConvolution except ImportError: @@ -77,7 +79,7 @@ def __init__(self, config, layer_idx: int, args=None): self.head_v_dim, eps=self.layer_norm_epsilon, activation=self.activation, - device=torch.cuda.current_device(), + device=accelerator.current_device(), dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(), ) @@ -159,6 +161,7 @@ def __init__( layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, + name: str | None = None, ): super().__init__( args, diff --git a/vime_plugins/models/qwen3_5_vl.py b/vime_plugins/models/qwen3_5_vl.py new file mode 100644 index 000000000..7460c7040 --- /dev/null +++ b/vime_plugins/models/qwen3_5_vl.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import torch +from megatron.core import mpu, tensor_parallel +from megatron.core.models.common.embeddings.rotary_pos_embedding import MultimodalRotaryEmbedding +from megatron.core.models.gpt import GPTModel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.module import MegatronModule +from transformers import AutoConfig + +from vime.utils import accelerator + +from .qwen3_5 import get_qwen3_5_spec +from .qwen3_5_vl_utils import build_packed_mrope_position_ids, gather_packed_input_ids, get_packed_cp_local_indices + + +class Qwen3_5MultimodalRotaryEmbedding(MultimodalRotaryEmbedding): + """Qwen3.5's interleaved temporal/height/width MRoPE layout.""" + + def forward( + self, + position_ids: torch.Tensor, + mrope_section: list[int], + packed_seq: bool = False, + cp_group=None, + ) -> torch.Tensor: + seq = position_ids.to(device=self.inv_freq.device, dtype=torch.float32) + if self.seq_len_interpolation_factor is not None: + seq *= 1 / self.seq_len_interpolation_factor + + inv_freq = self.inv_freq[None, None, :, None].expand(3, seq.shape[1], -1, 1) + freqs = (inv_freq @ seq[:, :, None, :]).transpose(2, 3) + + # Qwen3.5 interleaves T/H/W frequency bands instead of concatenating them. + mixed_freqs = freqs[0].clone() + for axis, offset in ((1, 1), (2, 2)): + mixed_freqs[..., offset : mrope_section[axis] * 3 : 3] = freqs[ + axis, ..., offset : mrope_section[axis] * 3 : 3 + ] + + emb = torch.cat((mixed_freqs, mixed_freqs), dim=-1)[..., None, :].transpose(0, 1).contiguous() + cp_group = cp_group or self.cp_group + if cp_group is not None and cp_group.size() > 1 and not packed_seq: + from megatron.core.models.common.embeddings.rope_utils import get_pos_emb_on_this_cp_rank + + emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) + return emb + + +def _load_vision_model(hf_config, dtype: torch.dtype, use_cpu_initialization: bool): + if hf_config.model_type == "qwen3_5_moe": + from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeVisionModel + + vision_model_cls = Qwen3_5MoeVisionModel + else: + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel + + vision_model_cls = Qwen3_5VisionModel + + device = torch.device("cpu") if use_cpu_initialization else accelerator.current_device() + with torch.device(device): + vision_model = vision_model_cls._from_config(hf_config.vision_config) + vision_model.to(dtype=dtype) + + # HF modules are replicated across TP ranks. Mark them explicitly so vime's + # direct weight exporter does not try to tensor-parallel all-gather them. + for parameter in vision_model.parameters(): + parameter.tensor_model_parallel = False + parameter.partition_dim = -1 + parameter.partition_stride = 1 + return vision_model + + +class Qwen3_5VLModel(MegatronModule): + """Megatron Qwen3.5 language model with a replicated Transformers ViT.""" + + def __init__( + self, + config, + language_layer_spec, + hf_config, + args, + *, + pre_process: bool, + post_process: bool, + vp_stage: int | None, + ) -> None: + super().__init__(config=config) + self.pre_process = pre_process + self.post_process = post_process + self.image_token_id = hf_config.image_token_id + self.video_token_id = hf_config.video_token_id + self.vision_start_token_id = hf_config.vision_start_token_id + self.spatial_merge_size = hf_config.vision_config.spatial_merge_size + + config.mrope_section = list(hf_config.text_config.rope_parameters["mrope_section"]) + config.apply_rope_fusion = False + + mtp_block_spec = None + if args.mtp_num_layers: + mtp_block_spec = get_gpt_mtp_block_spec( + config, + language_layer_spec, + use_transformer_engine=args.transformer_impl == "transformer_engine", + **({"vp_stage": vp_stage} if vp_stage is not None else {}), + ) + + self.language_model = GPTModel( + config=config, + transformer_layer_spec=language_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type="mrope", + rotary_percent=args.rotary_percent, + rotary_base=args.rotary_base, + rope_scaling=args.use_rope_scaling, + scatter_embedding_sequence_parallel=False, + mtp_block_spec=mtp_block_spec, + **({"vp_stage": vp_stage} if vp_stage is not None else {}), + ) + self.language_model.rotary_pos_emb = Qwen3_5MultimodalRotaryEmbedding( + kv_channels=config.kv_channels, + rotary_percent=args.rotary_percent, + rotary_interleaved=config.rotary_interleaved, + rotary_base=args.rotary_base, + ) + + self.model = torch.nn.Module() + self.model.visual = ( + _load_vision_model(hf_config, config.params_dtype, config.use_cpu_initialization) if pre_process else None + ) + self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights + + @property + def decoder(self): + return self.language_model.decoder + + def shared_embedding_or_output_weight(self): + return self.language_model.shared_embedding_or_output_weight() + + def set_input_tensor(self, input_tensor) -> None: + self.language_model.set_input_tensor(input_tensor) + + def _inject_vision_embeddings( + self, + input_ids: torch.Tensor, + full_input_ids: torch.Tensor, + cu_seqlens: torch.Tensor, + cp_group, + pixel_values: torch.Tensor | None, + pixel_values_videos: torch.Tensor | None, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + ) -> torch.Tensor: + embeddings = self.language_model.embedding(input_ids=input_ids, position_ids=None).clone() + embeddings_bsh = embeddings.transpose(0, 1).contiguous() + local_indices = get_packed_cp_local_indices( + cu_seqlens, + cp_group.size() if cp_group is not None else 1, + cp_group.rank() if cp_group is not None else 0, + input_ids.device, + ) + + for values, grids, token_id in ( + (pixel_values, image_grid_thw, self.image_token_id), + (pixel_values_videos, video_grid_thw, self.video_token_id), + ): + if values is None: + continue + if grids is None: + raise ValueError("Qwen3.5-VL pixel values require matching grid_thw") + vision_output = self.model.visual(values.to(dtype=self.model.visual.dtype), grid_thw=grids) + vision_embeddings = vision_output.pooler_output.to(device=embeddings.device, dtype=embeddings.dtype) + full_vision_positions = (full_input_ids[0] == token_id).nonzero(as_tuple=False).flatten() + if full_vision_positions.numel() != vision_embeddings.shape[0]: + raise ValueError( + f"Qwen3.5-VL token/features mismatch: {full_vision_positions.numel()} tokens, " + f"{vision_embeddings.shape[0]} features" + ) + + feature_indices = torch.full( + (full_input_ids.shape[1],), + -1, + dtype=torch.long, + device=input_ids.device, + ) + feature_indices[full_vision_positions] = torch.arange(vision_embeddings.shape[0], device=input_ids.device) + local_feature_indices = feature_indices[local_indices] + local_vision_mask = local_feature_indices >= 0 + if not torch.equal(local_vision_mask, input_ids[0] == token_id): + raise ValueError("Qwen3.5-VL CP token layout does not match its full packed sequence") + embeddings_bsh[0, local_vision_mask] = vision_embeddings[local_feature_indices[local_vision_mask]] + + embeddings = embeddings_bsh.transpose(0, 1).contiguous() + if self.config.sequence_parallel: + embeddings = tensor_parallel.scatter_to_sequence_parallel_region(embeddings).contiguous() + return embeddings + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + packed_seq_params: PackedSeqParams | None = None, + loss_mask: torch.Tensor | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if packed_seq_params is None: + raise ValueError("Qwen3.5-VL native training currently requires packed sequences") + + cp_group = mpu.get_context_parallel_group() + full_input_ids = gather_packed_input_ids(input_ids, packed_seq_params.cu_seqlens_q, cp_group) + + decoder_input = None + if self.pre_process: + decoder_input = self._inject_vision_embeddings( + input_ids, + full_input_ids, + packed_seq_params.cu_seqlens_q, + cp_group, + pixel_values, + pixel_values_videos, + image_grid_thw, + video_grid_thw, + ) + + if position_ids is None: + position_ids = build_packed_mrope_position_ids( + full_input_ids, + packed_seq_params.cu_seqlens_q, + image_grid_thw, + video_grid_thw, + image_token_id=self.image_token_id, + video_token_id=self.video_token_id, + vision_start_token_id=self.vision_start_token_id, + spatial_merge_size=self.spatial_merge_size, + ) + + return self.language_model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + labels=labels, + packed_seq_params=packed_seq_params, + loss_mask=loss_mask, + **kwargs, + ) + + +def get_qwen3_5_vl_model_provider(args, config, vp_stage): + """Return the native Qwen3.5-VL model provider.""" + + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + if not hasattr(hf_config, "vision_config"): + raise ValueError(f"{args.hf_checkpoint} is not a Qwen3.5-VL checkpoint") + language_layer_spec = get_qwen3_5_spec(args, config, vp_stage) + + def model_provider( + pre_process: bool = True, + post_process: bool = True, + vp_stage: int | None = None, + ) -> Qwen3_5VLModel: + return Qwen3_5VLModel( + config, + language_layer_spec, + hf_config, + args, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + + return model_provider diff --git a/vime_plugins/models/qwen3_5_vl_utils.py b/vime_plugins/models/qwen3_5_vl_utils.py new file mode 100644 index 000000000..26f59d1f7 --- /dev/null +++ b/vime_plugins/models/qwen3_5_vl_utils.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.distributed as dist + + +def get_packed_cp_local_indices( + cu_seqlens: Sequence[int] | torch.Tensor, + cp_size: int, + cp_rank: int, + device: torch.device, +) -> torch.Tensor: + """Map a THD CP rank's local tokens back to the full packed token stream.""" + + boundaries = [int(value) for value in cu_seqlens] + indices = [] + for start, end in zip(boundaries[:-1], boundaries[1:], strict=True): + sequence_length = end - start + if sequence_length % (2 * cp_size) != 0: + raise ValueError(f"Packed sequence length {sequence_length} must be divisible by 2 * CP size {cp_size}") + chunk_size = sequence_length // (2 * cp_size) + first = start + cp_rank * chunk_size + second = start + (2 * cp_size - cp_rank - 1) * chunk_size + indices.extend( + ( + torch.arange(first, first + chunk_size, device=device), + torch.arange(second, second + chunk_size, device=device), + ) + ) + return torch.cat(indices) if indices else torch.empty(0, dtype=torch.long, device=device) + + +def gather_packed_input_ids( + input_ids: torch.Tensor, + cu_seqlens: Sequence[int] | torch.Tensor, + cp_group, +) -> torch.Tensor: + """Reconstruct the full THD token stream from Megatron's two-chunk CP layout.""" + + if cp_group is None or cp_group.size() == 1: + return input_ids + + local_tokens = input_ids.flatten() + gathered_tokens = [torch.empty_like(local_tokens) for _ in range(cp_group.size())] + dist.all_gather(gathered_tokens, local_tokens, group=cp_group) + + full_tokens = torch.empty(int(cu_seqlens[-1]), dtype=input_ids.dtype, device=input_ids.device) + for rank, rank_tokens in enumerate(gathered_tokens): + indices = get_packed_cp_local_indices(cu_seqlens, cp_group.size(), rank, input_ids.device) + if indices.numel() != rank_tokens.numel(): + raise ValueError( + f"CP rank {rank} has {rank_tokens.numel()} tokens, expected {indices.numel()} from cu_seqlens" + ) + full_tokens[indices] = rank_tokens + return full_tokens.unsqueeze(0) + + +def _vision_positions( + start_position: int, + grid_thw: torch.Tensor, + spatial_merge_size: int, + device: torch.device, +) -> torch.Tensor: + grid_t, grid_h, grid_w = (int(value) for value in grid_thw.tolist()) + grid_h //= spatial_merge_size + grid_w //= spatial_merge_size + + temporal = torch.arange(grid_t, device=device).repeat_interleave(grid_h * grid_w) + start_position + height = torch.arange(grid_h, device=device).repeat_interleave(grid_w).repeat(grid_t) + start_position + width = torch.arange(grid_w, device=device).repeat(grid_h * grid_t) + start_position + return torch.stack((temporal, height, width)) + + +def build_packed_mrope_position_ids( + input_ids: torch.Tensor, + cu_seqlens: Sequence[int] | torch.Tensor, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + *, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + spatial_merge_size: int, +) -> torch.Tensor: + """Build Qwen3.5 MRoPE positions while resetting every packed sequence.""" + + if input_ids.shape[0] != 1: + raise ValueError(f"Packed Qwen3.5-VL input_ids must have batch size 1, got {tuple(input_ids.shape)}") + + device = input_ids.device + positions = torch.zeros((3, 1, input_ids.shape[1]), dtype=input_ids.dtype, device=device) + boundaries = [int(value) for value in cu_seqlens] + image_grids = iter(()) if image_grid_thw is None else iter(image_grid_thw) + + if video_grid_thw is None: + video_grids = iter(()) + else: + split_video_grids = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0).clone() + split_video_grids[:, 0] = 1 + video_grids = iter(split_video_grids) + + for start, end in zip(boundaries[:-1], boundaries[1:], strict=True): + tokens = input_ids[0, start:end] + token_list = tokens.tolist() + pieces = [] + cursor = 0 + current_position = 0 + + vision_starts = (tokens == vision_start_token_id).nonzero(as_tuple=False).flatten() + modality_tokens = tokens[vision_starts + 1] if vision_starts.numel() else () + + for modality_token in modality_tokens: + modality_token = int(modality_token) + if modality_token not in (image_token_id, video_token_id): + continue + try: + vision_start = token_list.index(modality_token, cursor) + grid = next(image_grids if modality_token == image_token_id else video_grids) + except (StopIteration, ValueError) as exc: + raise ValueError("Qwen3.5-VL tokens and vision grids do not match") from exc + + text_length = vision_start - cursor + if text_length: + pieces.append(torch.arange(text_length, device=device).view(1, -1).expand(3, -1) + current_position) + current_position += text_length + + vision_position_ids = _vision_positions(current_position, grid, spatial_merge_size, device) + pieces.append(vision_position_ids) + current_position += max(int(grid[1]), int(grid[2])) // spatial_merge_size + cursor = vision_start + vision_position_ids.shape[1] + + if cursor < len(token_list): + text_length = len(token_list) - cursor + pieces.append(torch.arange(text_length, device=device).view(1, -1).expand(3, -1) + current_position) + + if pieces: + sequence_positions = torch.cat(pieces, dim=1) + if sequence_positions.shape[1] != len(token_list): + raise ValueError( + "Qwen3.5-VL vision token count does not match its grid: " + f"expected {len(token_list)} positions, built {sequence_positions.shape[1]}" + ) + positions[:, 0, start:end] = sequence_positions + + try: + next(image_grids) + raise ValueError("Unused Qwen3.5-VL image grids") + except StopIteration: + pass + try: + next(video_grids) + raise ValueError("Unused Qwen3.5-VL video grids") + except StopIteration: + pass + return positions diff --git a/vime_plugins/models/qwen3_next.py b/vime_plugins/models/qwen3_next.py index 683cb2806..90b33f283 100644 --- a/vime_plugins/models/qwen3_next.py +++ b/vime_plugins/models/qwen3_next.py @@ -9,6 +9,8 @@ from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from transformers.activations import ACT2FN +from vime.utils import accelerator + from .hf_attention import _load_hf_config try: @@ -70,7 +72,7 @@ def __init__(self, config, layer_idx: int, args=None): self.head_v_dim, eps=self.layer_norm_epsilon, activation=self.activation, - device=torch.cuda.current_device(), + device=accelerator.current_device(), dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(), ) @@ -181,6 +183,7 @@ def __init__( layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, + name: str | None = None, ): super().__init__( args, diff --git a/vime_plugins/models/qwen3_omni_moe.py b/vime_plugins/models/qwen3_omni_moe.py new file mode 100644 index 000000000..9b0bbfbb9 --- /dev/null +++ b/vime_plugins/models/qwen3_omni_moe.py @@ -0,0 +1,981 @@ +"""Native Megatron model for Qwen3-Omni-MoE Thinker training. + +Architecture (Thinker-only, Talker/Code2Wav are frozen and not trained): + HF audio encoder (Qwen3OmniMoeAudioEncoder, replicated on first PP stage) + HF vision encoder (Qwen3OmniMoeVisionEncoder, replicated on first PP stage) + + Megatron GPTModel (MoE language model with M-RoPE, deepstack) + +The forward pass: + 1. Computes text embeddings from `input_ids`. + 2. Runs the HF vision encoder on `pixel_values`+`image_grid_thw` + (and `pixel_values_videos`+`video_grid_thw` if present), scatters the + resulting vision embeddings into the combined embedding tensor at + positions where `input_ids == image_token_id` / `video_token_id`. + 3. Runs the HF audio encoder on `input_features`+`feature_attention_mask`, + scatters the resulting audio embeddings at positions where + `input_ids == audio_token_id`. + 4. Computes 3D M-RoPE position IDs from the full input_ids + grid info + (audio-aware, ported from Relax's get_rope_index). + 5. Forwards the combined embeddings + M-RoPE position IDs through the + Megatron GPTModel (MoE) language model. +""" + +from __future__ import annotations + +import logging +from copy import deepcopy + +import torch +from megatron.core import InferenceParams, mpu, parallel_state, tensor_parallel +from megatron.core.models.gpt import GPTModel as MCoreGPTModel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.module import MegatronModule +from megatron.core.utils import deprecate_inference_params + +from .qwen3_5_vl import Qwen3_5MultimodalRotaryEmbedding +from .qwen3_5_vl_utils import gather_packed_input_ids +from .qwen3_omni_transformer import Qwen3OmniTransformerBlock, split_deepstack_embeddings + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# THD <-> batch-sequence helpers +# --------------------------------------------------------------------------- +def _thd_to_batch_seq(packed: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: + """Unpack THD-format [1, T, ...] to [bs, max_seq, ...] using cu_seqlens.""" + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + max_seq = seqlens.max().item() + bs = len(cu_seqlens) - 1 + out = packed.new_zeros(bs, max_seq, *packed.shape[2:]) + for i, sl in enumerate(seqlens): + out[i, :sl] = packed[0, cu_seqlens[i] : cu_seqlens[i] + sl] + return out + + +def _batch_seq_to_thd(unpacked: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: + """Pack [bs, max_seq, ...] back to THD [1, T, ...].""" + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + total = cu_seqlens[-1].item() + out = unpacked.new_zeros(1, total, *unpacked.shape[2:]) + for i, sl in enumerate(seqlens): + out[0, cu_seqlens[i] : cu_seqlens[i] + sl] = unpacked[i, :sl] + return out + + +def _gather_input_ids_from_cp( + input_ids: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> torch.Tensor: + """Reconstruct the full packed sequence from Megatron's zigzag CP layout.""" + return gather_packed_input_ids(input_ids, cu_seqlens, parallel_state.get_context_parallel_group()) + + +# --------------------------------------------------------------------------- +# Audio-aware M-RoPE position ID computation (ported from Relax) +# --------------------------------------------------------------------------- +def _get_feat_extract_output_lengths(input_lengths): + """Computes the output length of the conv layers and the audio encoder.""" + input_lengths_leave = input_lengths % 100 + feat_lengths = (input_lengths_leave - 1) // 2 + 1 + output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + return output_lengths + + +def _get_rope_index( + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + audio_token_id: int, + vision_start_token_id: int, + audio_start_token_id: int, + input_ids: torch.Tensor, + image_grid_thw: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + audio_seqlens: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + use_audio_in_video: bool = False, + second_per_grids: torch.Tensor | None = None, + position_id_per_seconds: int = 1, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate RoPE position indices for multimodal inputs (audio+image+video). + + Ported from relax.models.qwen_omni.modeling_qwen3_omni.utils.get_rope_index. + Returns position_ids of shape [3, batch, seq] for M-RoPE. + """ + # Do NOT split video_grid_thw by repeat_interleave. + # The Qwen3-Omni processor does NOT insert timestamp tokens between video + # frames for pure video (use_audio_in_video=False). Input_ids have ONE + # + (grid_t*grid_h*grid_w/merge^2) + . + # Splitting grid_t into t=1 entries would under-count video tokens and break + # M-RoPE vs vLLM (see vLLM get_mrope_input_positions). + + mrope_position_deltas = [] + if image_grid_thw is not None or video_grid_thw is not None or audio_seqlens is not None: + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=torch.float, + device=input_ids.device, + ) + image_index, video_index, audio_index = 0, 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids_i in enumerate(total_input_ids): + input_ids_i = input_ids_i[attention_mask[i] == 1] + + vision_start_indices = torch.argwhere(input_ids_i == vision_start_token_id).squeeze(1) + vision_tokens = input_ids_i[vision_start_indices + 1] + audio_nums = torch.sum(input_ids_i == audio_start_token_id) + image_nums = (vision_tokens == image_token_id).sum() + video_nums = ( + (vision_tokens == audio_start_token_id).sum() + if use_audio_in_video + else (vision_tokens == video_token_id).sum() + ) + input_tokens = input_ids_i.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums + multimodal_nums = image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums + + for _ in range(multimodal_nums): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + if (image_token_id in input_tokens or video_token_id in input_tokens) and ( + remain_videos > 0 or remain_images > 0 + ): + ed_vision_start = input_tokens.index(vision_start_token_id, st) + else: + ed_vision_start = len(input_tokens) + 1 + if audio_token_id in input_tokens and remain_audios > 0: + ed_audio_start = input_tokens.index(audio_start_token_id, st) + else: + ed_audio_start = len(input_tokens) + 1 + min_ed = min(ed_vision_start, ed_audio_start) + + # ---------- text ---------- + text_len = min_ed - st + if text_len > 0: + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + st_idx += text_len + + # ---------- BOS ---------- + if min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start: + bos_len, eos_len = 2, 2 + else: + bos_len, eos_len = 1, 1 + + llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx) + st_idx += bos_len + + # Audio Only + if min_ed == ed_audio_start: + audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_index]) + llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx + llm_pos_ids_list.append(llm_pos_ids) + + st += text_len + bos_len + audio_len + eos_len + audio_index += 1 + remain_audios -= 1 + + # Image Only + elif min_ed == ed_vision_start and input_ids_i[ed_vision_start + 1] == image_token_id: + t, h, w = ( + image_grid_thw[image_index][0].item(), + image_grid_thw[image_index][1].item(), + image_grid_thw[image_index][2].item(), + ) + t_index = (torch.arange(t) * 1 * position_id_per_seconds).float() + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + h_index = ( + torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float() + ) + w_index = ( + torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float() + ) + t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float() + _llm_pos_ids = torch.stack([t_index, h_index, w_index]) + llm_pos_ids_list.append(_llm_pos_ids + st_idx) + + image_len = image_grid_thw[image_index].prod().item() // (spatial_merge_size**2) + st += int(text_len + bos_len + image_len + eos_len) + image_index += 1 + remain_images -= 1 + + # Video Only + elif min_ed == ed_vision_start and input_ids_i[ed_vision_start + 1] == video_token_id: + t, h, w = ( + video_grid_thw[video_index][0].item(), + video_grid_thw[video_index][1].item(), + video_grid_thw[video_index][2].item(), + ) + t_index = ( + torch.arange(t) * second_per_grids[video_index].cpu().float() * position_id_per_seconds + ).float() + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + h_index = ( + torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float() + ) + w_index = ( + torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float() + ) + t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float() + _llm_pos_ids = torch.stack([t_index, h_index, w_index]) + llm_pos_ids_list.append(_llm_pos_ids + st_idx) + + video_len = video_grid_thw[video_index].prod().item() // (spatial_merge_size**2) + st += int(text_len + bos_len + video_len + eos_len) + video_index += 1 + remain_videos -= 1 + + # Audio in Video + elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start: + audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_index]) + audio_llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx + + t, h, w = ( + video_grid_thw[video_index][0].item(), + video_grid_thw[video_index][1].item(), + video_grid_thw[video_index][2].item(), + ) + t_index = ( + torch.arange(t) * second_per_grids[video_index].cpu().float() * position_id_per_seconds + ).float() + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + h_index = ( + torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float() + ) + w_index = ( + torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float() + ) + t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float() + _llm_pos_ids = torch.stack([t_index, h_index, w_index]) + llm_pos_ids_list_temp = [_llm_pos_ids + st_idx] + video_llm_pos_ids = torch.cat(llm_pos_ids_list_temp, dim=1) + + video_data_index, audio_data_index = 0, 0 + while ( + video_data_index < video_llm_pos_ids.shape[-1] + and audio_data_index < audio_llm_pos_ids.shape[-1] + ): + if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]: + llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1]) + video_data_index += 1 + else: + llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1]) + audio_data_index += 1 + if video_data_index < video_llm_pos_ids.shape[-1]: + llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]]) + if audio_data_index < audio_llm_pos_ids.shape[-1]: + llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]]) + video_len = video_grid_thw[video_index].prod().item() // (spatial_merge_size**2) + + st += int(text_len + bos_len + audio_len + video_len + eos_len) + audio_index += 1 + video_index += 1 + remain_videos -= 1 + remain_audios -= 1 + else: + raise RuntimeError("unexpected error in get_rope_index") + + # ---------- EOS ---------- + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx) + + # tail text + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat([item.float() for item in llm_pos_ids_list], dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids_i)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + # fallback (pure text) + if attention_mask is not None: + position_ids = attention_mask.float().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True) + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + +# --------------------------------------------------------------------------- +# GPTModel with DeepStack support (keeps MCore mrope, swaps decoder only) +# --------------------------------------------------------------------------- +class Qwen3OmniMultimodalRotaryEmbedding(Qwen3_5MultimodalRotaryEmbedding): + """Qwen3 interleaved MRoPE with packed-sequence CP handling.""" + + def __init__(self, *args, cp_group=None, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.cp_group = cp_group + self.is_thd_format = False + + def forward(self, position_ids, mrope_section, packed_seq_params=None, **kwargs): + return super().forward( + position_ids, + mrope_section, + packed_seq=self.is_thd_format, + cp_group=self.cp_group, + ) + + +class Qwen3OmniMoeGPTModel(MCoreGPTModel): + """Qwen3-Omni GPT model with DeepStack support. + + Inherits GPTModel to keep MCore's MultimodalRotaryEmbedding (proven for text + training). Only replaces the decoder to add DeepStack + injection at the first N decoder layers. + """ + + def __init__( + self, + config, + transformer_layer_spec, + vocab_size: int, + max_sequence_length: int, + pre_process: bool = True, + post_process: bool = True, + fp16_lm_cross_entropy: bool = False, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = False, + position_embedding_type: str = "learned_absolute", + rotary_percent: float = 1.0, + rotary_base: int = 10000, + rope_scaling: bool = False, + rope_scaling_factor: float = 8.0, + scatter_embedding_sequence_parallel: bool = True, + seq_len_interpolation_factor=None, + mtp_block_spec=None, + vp_stage=None, + pg_collection=None, + ) -> None: + super().__init__( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=vocab_size, + max_sequence_length=max_sequence_length, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=fp16_lm_cross_entropy, + parallel_output=parallel_output, + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + position_embedding_type=position_embedding_type, + rotary_percent=rotary_percent, + rotary_base=rotary_base, + rope_scaling=rope_scaling, + rope_scaling_factor=rope_scaling_factor, + scatter_embedding_sequence_parallel=scatter_embedding_sequence_parallel, + seq_len_interpolation_factor=seq_len_interpolation_factor, + mtp_block_spec=mtp_block_spec, + vp_stage=vp_stage, + pg_collection=pg_collection, + ) + # Match HF/vLLM's interleaved multimodal RoPE layout. + # CRITICAL: MCore's MultimodalRotaryEmbedding uses NON-interleaved mrope + # layout [T48,H40,W40,...] which diverges from vLLM/HF interleaved layout + # [T24,H24,W24,...] when t!=h!=w (video). For text (t=h=w) both are identical. + cp_group = None + if pg_collection is not None and getattr(pg_collection, "cp", None) is not None: + cp_group = pg_collection.cp + else: + from megatron.core import parallel_state + + cp_group = parallel_state.get_context_parallel_group(check_initialized=False) + self.rotary_pos_emb = Qwen3OmniMultimodalRotaryEmbedding( + kv_channels=self.config.kv_channels, + rotary_percent=rotary_percent, + rotary_interleaved=False, # bridge asserts not interleaved; uses apply_interleaved_mrope internally + seq_len_interpolation_factor=seq_len_interpolation_factor, + rotary_base=rotary_base, + cp_group=cp_group, + ) + # Rebuild the decoder with DeepStack injection. + self.decoder = Qwen3OmniTransformerBlock( + config=self.config, + spec=transformer_layer_spec, + pre_process=self.pre_process, + post_process=self.post_process, + vp_stage=vp_stage, + pg_collection=pg_collection, + ) + + def forward( + self, + input_ids, + position_ids, + attention_mask, + decoder_input=None, + labels=None, + inference_context=None, + packed_seq_params=None, + extra_block_kwargs=None, + runtime_gather_output=None, + *, + inference_params=None, + loss_mask=None, + # args for deepstack + visual_pos_masks=None, + deepstack_visual_embeds=None, + ): + """Forward pass with DeepStack visual embedding injection.""" + inference_context = deprecate_inference_params(inference_context, inference_params) + + preproc_output = self._preprocess( + input_ids=input_ids, + position_ids=position_ids, + decoder_input=decoder_input, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + ( + decoder_input, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + ) = preproc_output[:5] + + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_context=inference_context, + 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, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + **(extra_block_kwargs or {}), + ) + + result = self._postprocess( + hidden_states=hidden_states, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + mtp_in_postprocess=self.mtp_process, + loss_mask=loss_mask, + decoder_input=decoder_input, + attention_mask=attention_mask, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + runtime_gather_output=runtime_gather_output, + extra_block_kwargs=extra_block_kwargs, + inference_context=inference_context, + ) + return result + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- +class Qwen3OmniMoeVLModel(MegatronModule): + """Qwen3-Omni-MoE Thinker model for Megatron training. + + Wraps an HF audio encoder and an HF vision encoder (only on first PP stage) + together with a standard Megatron Core GPTModel configured for M-RoPE + (MoE language model). + + Thinker-only training: the Talker and Code2Wav modules are not loaded. + The audio and vision encoders are frozen by default (RL only trains the + language model). + """ + + def __init__( + self, + language_transformer_config, + language_transformer_layer_spec, + hf_audio_config, + hf_vision_config, + parallel_output: bool = True, + pre_process: bool = True, + post_process: bool = True, + pg_collection=None, + ) -> None: + super().__init__(config=language_transformer_config) + + self.pre_process = pre_process + self.post_process = post_process + self.pg_collection = pg_collection + + self.image_token_id = language_transformer_config.image_token_id + self.video_token_id = language_transformer_config.video_token_id + self.vision_start_token_id = language_transformer_config.vision_start_token_id + self.audio_token_id = language_transformer_config.audio_token_id + self.audio_start_token_id = language_transformer_config.audio_start_token_id + self.spatial_merge_size = language_transformer_config.spatial_merge_size + self.position_id_per_seconds = language_transformer_config.position_id_per_seconds + self.use_audio_in_video = getattr(language_transformer_config, "use_audio_in_video", False) + + self.share_embeddings_and_output_weights = False + + # Encoders -- only on the first pipeline stage + self.audio_model = None + self.vision_model = None + if self.pre_process: + from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeAudioEncoder, + Qwen3OmniMoeVisionEncoder, + ) + + self.audio_model = Qwen3OmniMoeAudioEncoder._from_config(hf_audio_config) + self.vision_model = Qwen3OmniMoeVisionEncoder._from_config(hf_vision_config) + # Freeze encoders -- not trained during RL + self.audio_model.requires_grad_(False) + self.audio_model.eval() + self.vision_model.requires_grad_(False) + self.vision_model.eval() + + for parameter in (*self.audio_model.parameters(), *self.vision_model.parameters()): + parameter.tensor_model_parallel = False + parameter.partition_dim = -1 + parameter.partition_stride = 1 + if torch.cuda.is_available(): + # Keep encoder param dtype (often bf16 from HF config); only move device. + _enc_device = torch.device(f"cuda:{torch.cuda.current_device()}") + _audio_dtype = next(self.audio_model.parameters()).dtype + _vision_dtype = next(self.vision_model.parameters()).dtype + self.audio_model = self.audio_model.to(device=_enc_device, dtype=_audio_dtype) + self.vision_model = self.vision_model.to(device=_enc_device, dtype=_vision_dtype) + + # Language model -- Megatron GPT with M-RoPE + DeepStack support + self.language_model = Qwen3OmniMoeGPTModel( + config=language_transformer_config, + transformer_layer_spec=language_transformer_layer_spec, + vocab_size=language_transformer_config.vocab_size, + max_sequence_length=language_transformer_config.language_max_sequence_length, + parallel_output=parallel_output, + position_embedding_type="mrope", + rotary_percent=language_transformer_config.rotary_percent, + pre_process=self.pre_process, + post_process=self.post_process, + rotary_base=language_transformer_config.rotary_base, + fp16_lm_cross_entropy=language_transformer_config.fp16_lm_cross_entropy, + share_embeddings_and_output_weights=language_transformer_config.share_embeddings_and_output_weights, + scatter_embedding_sequence_parallel=False, + pg_collection=pg_collection, + ) + + self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights + + # -- helpers required by Megatron pipeline engine ----------------------- + + def shared_embedding_or_output_weight(self): + return self.language_model.shared_embedding_or_output_weight() + + def set_input_tensor(self, input_tensor): + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + assert len(input_tensor) == 1 + if self.pre_process: + self.encoder_hidden_state = input_tensor[0] + else: + self.language_model.set_input_tensor(input_tensor[0]) + + # -- encoder helpers ---------------------------------------------------- + + def _get_vision_features(self, pixel_values, image_grid_thw): + pixel_values = pixel_values.to(dtype=self.vision_model.dtype) + with torch.no_grad(): + outputs = self.vision_model(pixel_values, grid_thw=image_grid_thw) + if hasattr(outputs, "pooler_output"): + return outputs.pooler_output, outputs.deepstack_features + return outputs + + def _get_audio_features(self, input_features, feature_lens): + with torch.no_grad(): + outputs = self.audio_model( + input_features.to(dtype=self.audio_model.dtype), + feature_lens=feature_lens, + ) + return outputs.last_hidden_state + + # -- forward ------------------------------------------------------------ + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor = None, + attention_mask: torch.Tensor = None, + labels: torch.Tensor = None, + loss_mask: torch.Tensor = None, + inference_params: InferenceParams = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + # multimodal kwargs + pixel_values: torch.Tensor = None, + image_grid_thw: torch.Tensor = None, + pixel_values_videos: torch.Tensor = None, + video_grid_thw: torch.Tensor = None, + image_input_mask: torch.Tensor = None, + video_second_per_grid: torch.Tensor = None, + # audio kwargs + input_features: torch.Tensor = None, + feature_attention_mask: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass of the Qwen3-Omni Thinker model. + + Args: + input_ids: [batch, seq] or THD [1, T] text token ids. + position_ids: optional, otherwise computed from input_ids. + attention_mask: text attention mask. + pixel_values: image pixel values (flat, [N_pix, C*P*P]). + image_grid_thw: [num_images, 3] (T, H, W) per image. + pixel_values_videos: video pixel values. + video_grid_thw: [num_videos, 3] (T, H, W) per video. + image_input_mask: optional precomputed image mask. + video_second_per_grid: seconds per video grid (for MRoPE t_index). + input_features: audio mel features [batch, channels, time]. + feature_attention_mask: audio attention mask [batch, time]. + """ + assert inference_params is None, "Inference not supported" + + # Extract cu_seqlens and CP info early + cu_seqlens = None + if packed_seq_params is not None: + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + cp_size = parallel_state.get_context_parallel_world_size() + + # Audio feature lengths (computed once, used by both audio encoder and MRoPE) + audio_feature_lengths = None + if feature_attention_mask is not None: + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + + # Vision bookkeeping + video_start_index = 0 + vision_grid_thw = None + vision_data = None + image_mask = None + video_mask = None + deepstack_feature_lists = None + + combined_embeddings = None + visual_pos_masks = None + + if self.pre_process: + # ========================= + # Vision (image / video) + # ========================= + if image_grid_thw is not None or video_grid_thw is not None: + if image_grid_thw is not None: + image_mask = image_input_mask + if image_mask is None: + image_mask = (input_ids == self.image_token_id).contiguous() + vision_grid_thw = image_grid_thw + vision_data = pixel_values + video_start_index = image_mask.sum().item() + else: + video_start_index = 0 + + if video_grid_thw is not None: + video_mask = (input_ids == self.video_token_id).contiguous() + if vision_grid_thw is not None: + vision_grid_thw = torch.cat([vision_grid_thw, video_grid_thw], dim=0) + vision_data = torch.cat([vision_data, pixel_values_videos], dim=0) + else: + vision_grid_thw = video_grid_thw + vision_data = pixel_values_videos + + vision_embeds = None + if vision_grid_thw is not None and vision_grid_thw.shape[0] > 0: + vision_embeds, deepstack_feature_lists = self._get_vision_features(vision_data, vision_grid_thw) + vision_embeds = vision_embeds.to(dtype=self.language_model.embedding.word_embeddings.weight.dtype) + + # ========================= + # Text embeddings + # ========================= + combined_embeddings = self.language_model.embedding( + input_ids=input_ids, + position_ids=None, + ).clone() # [seq, batch, hidden] + + # ========================= + # Scatter vision embeds + # ========================= + if vision_embeds is not None: + if video_start_index == 0: + image_embeds = None + video_embeds = vision_embeds + elif video_start_index == vision_embeds.shape[0]: + image_embeds = vision_embeds + video_embeds = None + elif 0 < video_start_index < vision_embeds.shape[0]: + image_embeds = vision_embeds[:video_start_index] + video_embeds = vision_embeds[video_start_index:] + else: + raise ValueError( + f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got " + f"{video_start_index}" + ) + + # [seq, bs, h] -> [bs, seq, h] for masked scatter + combined_embeddings_bsh = combined_embeddings.transpose(0, 1).contiguous() + if image_embeds is not None: + combined_embeddings_bsh[image_mask] = image_embeds + if video_embeds is not None: + combined_embeddings_bsh[video_mask] = video_embeds + combined_embeddings = combined_embeddings_bsh.transpose(0, 1).contiguous() + + if image_embeds is not None and video_embeds is not None: + visual_pos_masks = image_mask | video_mask + elif image_embeds is not None: + visual_pos_masks = image_mask + elif video_embeds is not None: + visual_pos_masks = video_mask + + # ========================= + # Audio + # ========================= + if input_features is not None: + audio_mask = (input_ids == self.audio_token_id).contiguous() + # Flatten input_features using feature_attention_mask + if feature_attention_mask is not None: + input_features_flat = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0) + else: + input_features_flat = input_features + + feature_lens = ( + audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1) + ) + + audio_embeds = self._get_audio_features(input_features_flat, feature_lens) + audio_embeds = audio_embeds.to(combined_embeddings.dtype) + + combined_embeddings_bsh = combined_embeddings.transpose(0, 1).contiguous() + combined_embeddings_bsh[audio_mask] = audio_embeds + combined_embeddings = combined_embeddings_bsh.transpose(0, 1).contiguous() + + # Scatter to sequence-parallel region if needed + if self.config.sequence_parallel: + combined_embeddings = tensor_parallel.scatter_to_sequence_parallel_region(combined_embeddings) + combined_embeddings = combined_embeddings.contiguous() + + # ========================= + # Compute M-RoPE position IDs + # ========================= + # position_ids must be available on ALL PP stages for rotary embeddings. + pp_size = parallel_state.get_pipeline_model_parallel_world_size() + + if position_ids is None: + if self.pre_process: + # Reconstruct full input_ids if CP > 1 + if cu_seqlens is not None: + if cp_size > 1: + full_input_ids = _gather_input_ids_from_cp(input_ids, cu_seqlens) + else: + full_input_ids = input_ids + input_ids_batch_seq = _thd_to_batch_seq(full_input_ids, cu_seqlens) + else: + input_ids_batch_seq = input_ids + + # If no multimodal inputs at all, fall back to pure-text positions + has_multimodal = ( + (image_grid_thw is not None and image_grid_thw.numel() > 0) + or (video_grid_thw is not None and video_grid_thw.numel() > 0) + or audio_feature_lengths is not None + ) + + if has_multimodal: + pos_batch_seq, _ = _get_rope_index( + spatial_merge_size=self.spatial_merge_size, + image_token_id=self.image_token_id, + video_token_id=self.video_token_id, + audio_token_id=self.audio_token_id, + vision_start_token_id=self.vision_start_token_id, + audio_start_token_id=self.audio_start_token_id, + input_ids=input_ids_batch_seq, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + audio_seqlens=audio_feature_lengths, + attention_mask=None, + use_audio_in_video=self.use_audio_in_video, + second_per_grids=video_second_per_grid, + position_id_per_seconds=self.position_id_per_seconds, + ) + else: + # Pure text: standard 1D positions replicated across 3 dims + bs, seq_len = input_ids_batch_seq.shape + pos = torch.arange(seq_len, device=input_ids_batch_seq.device).unsqueeze(0).expand(bs, -1) + pos_batch_seq = torch.stack([pos, pos, pos], dim=0) # [3, bs, seq] + + if cu_seqlens is not None: + pos_packed = _batch_seq_to_thd(pos_batch_seq.permute(1, 2, 0), cu_seqlens) + position_ids = pos_packed.permute(2, 0, 1).contiguous() # [3, 1, T_global] + else: + position_ids = pos_batch_seq # [3, bs, seq] + else: + # Non-first PP stage: allocate buffer with correct shape + if cu_seqlens is not None: + T = cu_seqlens[-1].item() + position_ids = torch.zeros(3, 1, T, dtype=torch.float, device=torch.cuda.current_device()) + else: + raise NotImplementedError( + "Non-THD position_ids broadcast not yet supported for non-first PP stages" + ) + + # Broadcast position_ids from first to all PP stages + if pp_size > 1: + src = parallel_state.get_pipeline_model_parallel_first_rank() + torch.distributed.broadcast( + position_ids, + src=src, + group=parallel_state.get_pipeline_model_parallel_group(), + ) + + # ========================= + # Split deepstack features for SP / CP + # ========================= + if self.config.sequence_parallel and visual_pos_masks is not None and deepstack_feature_lists is not None: + if self.pg_collection is not None: + tp_size = self.pg_collection.tp.size() + tp_rank = self.pg_collection.tp.rank() + else: + tp_size = mpu.get_tensor_model_parallel_world_size() + tp_rank = mpu.get_tensor_model_parallel_rank() + visual_pos_masks, deepstack_feature_lists = split_deepstack_embeddings( + visual_pos_masks, + deepstack_feature_lists, + tp_size=tp_size, + tp_rank=tp_rank, + sequence_parallel=True, + ) + + # ========================= + # Packed THD RoPE is sliced by attention, not by the embedding module. + # ========================= + # Standard Qwen3-VL model sets is_thd_format=True dynamically (model.py:805,822) + # when using packed sequences with CP. Otherwise the embedding module + # slices emb along CP (because + # packed_seq kwarg is swallowed by **kwargs and is_thd_format stays False), + # producing freqs with T_global/cp_size entries. Then _apply_rotary_pos_emb_thd + # CASE 2 (_get_thd_freqs_on_this_cp_rank) accesses out-of-bounds indices for + # long sequences, producing a shorter freqs_packed that mismatches t. + # Fix: set is_thd_format=True for packed (THD) sequences so CP slicing is + # skipped here; _apply_rotary_pos_emb_thd handles CP per-sequence internally. + if hasattr(self.language_model, "rotary_pos_emb") and hasattr( + self.language_model.rotary_pos_emb, "is_thd_format" + ): + self.language_model.rotary_pos_emb.is_thd_format = cu_seqlens is not None + + # ========================= + # Language model forward + # ========================= + # NOTE: visual_pos_masks and deepstack_visual_embeds are Qwen3-Omni-specific + # DeepStack parameters. Standard Megatron GPTModel does not accept them; they + # require custom decoder layers. Only pass when not None so text-only training + # works with the standard GPTModel. Visual inputs need custom decoder support. + language_model_kwargs = dict( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=combined_embeddings, + labels=labels, + loss_mask=loss_mask, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + ) + if visual_pos_masks is not None: + language_model_kwargs["visual_pos_masks"] = visual_pos_masks + if deepstack_feature_lists is not None: + language_model_kwargs["deepstack_visual_embeds"] = deepstack_feature_lists + if extra_block_kwargs: + language_model_kwargs.update(extra_block_kwargs) + + output = self.language_model(**language_model_kwargs) + + return output + + +# --------------------------------------------------------------------------- +# Native model provider +# --------------------------------------------------------------------------- +def get_qwen3_omni_model_provider(args, config, vp_stage): + """Return the native Qwen3-Omni Thinker provider.""" + + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + thinker_config = hf_config.thinker_config + audio_config = deepcopy(thinker_config.audio_config) + vision_config = deepcopy(thinker_config.vision_config) + audio_config.torch_dtype = config.params_dtype + vision_config.torch_dtype = config.params_dtype + + text_config = thinker_config.text_config + rope_config = getattr(text_config, "rope_parameters", None) or getattr(text_config, "rope_scaling", None) or {} + values = { + "audio_start_token_id": thinker_config.audio_start_token_id, + "audio_token_id": thinker_config.audio_token_id, + "fp16_lm_cross_entropy": args.fp16_lm_cross_entropy, + "image_token_id": thinker_config.image_token_id, + "language_max_sequence_length": args.max_position_embeddings, + "mrope_section": list(rope_config.get("mrope_section", [24, 20, 20])), + "position_id_per_seconds": thinker_config.position_id_per_seconds, + "rotary_base": args.rotary_base, + "rotary_percent": args.rotary_percent, + "share_embeddings_and_output_weights": not args.untie_embeddings_and_output_weights, + "spatial_merge_size": vision_config.spatial_merge_size, + "use_audio_in_video": getattr(thinker_config, "use_audio_in_video", False), + "video_token_id": thinker_config.video_token_id, + "vision_start_token_id": thinker_config.vision_start_token_id, + "vocab_size": args.padded_vocab_size, + } + for name, value in values.items(): + setattr(config, name, value) + + layer_spec = get_gpt_layer_with_transformer_engine_spec( + num_experts=args.num_experts, + moe_grouped_gemm=args.moe_grouped_gemm, + qk_layernorm=args.qk_layernorm, + fp8=False, + ) + + def model_provider( + pre_process: bool = True, + post_process: bool = True, + vp_stage: int | None = None, + ) -> Qwen3OmniMoeVLModel: + return Qwen3OmniMoeVLModel( + language_transformer_config=config, + language_transformer_layer_spec=layer_spec, + hf_audio_config=audio_config, + hf_vision_config=vision_config, + pre_process=pre_process, + post_process=post_process, + ) + + return model_provider diff --git a/vime_plugins/models/qwen3_omni_transformer.py b/vime_plugins/models/qwen3_omni_transformer.py new file mode 100644 index 000000000..8e8576a95 --- /dev/null +++ b/vime_plugins/models/qwen3_omni_transformer.py @@ -0,0 +1,381 @@ +from __future__ import annotations + +from contextlib import nullcontext + +import torch +from megatron.core import tensor_parallel +from megatron.core.enums import Fp8Recipe +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import TransformerBlock, TransformerBlockSubmodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor +from torch import Tensor + +try: + from megatron.core.extensions.transformer_engine import te_checkpoint +except ImportError: + te_checkpoint = None + + +def split_deepstack_embeddings( + visual_mask: torch.Tensor, + embeddings: list[torch.Tensor], + *, + tp_size: int, + tp_rank: int, + sequence_parallel: bool, +) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Select the sequence-parallel slice of DeepStack visual features.""" + if not sequence_parallel or tp_size == 1: + return visual_mask, embeddings + if visual_mask.shape[-1] % tp_size: + raise ValueError(f"DeepStack sequence length must be divisible by TP size {tp_size}") + + mask_chunks = visual_mask.chunk(tp_size, dim=-1) + lengths = torch.stack([chunk.sum(-1) for chunk in mask_chunks], dim=-1) + offsets = torch.cat((lengths.new_zeros(1), lengths.flatten().cumsum(0))).tolist() + slices = [ + slice(offsets[batch * tp_size + tp_rank], offsets[batch * tp_size + tp_rank + 1]) + for batch in range(visual_mask.shape[0]) + ] + local_embeddings = [torch.cat([embedding[index] for index in slices]) for embedding in embeddings] + return mask_chunks[tp_rank], local_embeddings + + +class Qwen3OmniTransformerBlock(TransformerBlock): + """Transformer block that injects Qwen3-Omni DeepStack features.""" + + def __init__( + self, + config: TransformerConfig, + spec: TransformerBlockSubmodules | ModuleSpec, + post_layer_norm: bool = True, + pre_process: bool = True, + post_process: bool = True, + vp_stage: int | None = None, + pg_collection: ProcessGroupCollection | None = None, + ): + super().__init__( + config=config, + spec=spec, + post_layer_norm=post_layer_norm, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + pg_collection=pg_collection, + ) + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + self.pg_collection = pg_collection + self.cp_group = pg_collection.cp + self.tp_group = pg_collection.tp + self.pp_group = pg_collection.pp + + def _checkpointed_forward( + self, + hidden_states: Tensor, + attention_mask: Tensor, + context: Tensor, + context_mask: Tensor, + rotary_pos_emb: Tensor, + attention_bias: Tensor, + packed_seq_params: PackedSeqParams, + use_inner_fp8_context: bool, + # args for deepstack + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + ): + """Forward method with activation checkpointing.""" + + def custom(start: int, end: int): + def custom_forward( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + visual_pos_masks, + *deepstack_visual_embeds_args, + ): + deepstack_visual_embeds = list(deepstack_visual_embeds_args) if deepstack_visual_embeds_args else None + for index in range(start, end): + layer = self._get_layer(index) + inner_fp8_context = ( + get_fp8_context(self.config, layer.layer_number - 1) + if use_inner_fp8_context + else nullcontext() + ) + with inner_fp8_context: + hidden_states, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + attention_bias=attention_bias, + inference_context=None, + packed_seq_params=packed_seq_params, + ) + + if self.pre_process and deepstack_visual_embeds is not None: + l_no = layer.layer_number - 1 + if l_no in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[l_no], + ) + return hidden_states, context + + return custom_forward + + deepstack_visual_embeds_tuple = tuple(deepstack_visual_embeds) if deepstack_visual_embeds else () + + def checkpoint_handler(forward_func): + """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`""" + if self.config.fp8: + return te_checkpoint( + forward_func, + self.config.distribute_saved_activations, + tensor_parallel.random.get_cuda_rng_tracker, + self.tp_group, + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + visual_pos_masks, + *deepstack_visual_embeds_tuple, + ) + else: + return tensor_parallel.checkpoint( + forward_func, + self.config.distribute_saved_activations, + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + visual_pos_masks, + *deepstack_visual_embeds_tuple, + ) + + if self.config.recompute_method == "uniform": + # Uniformly divide the total number of Transformer layers and checkpoint + # the input activation of each divided chunk. + # A method to further reduce memory usage reducing checkpoints. + layer_idx = 0 + while layer_idx < self.num_layers_per_pipeline_rank: + hidden_states, context = checkpoint_handler( + custom(layer_idx, layer_idx + self.config.recompute_num_layers) + ) + + layer_idx += self.config.recompute_num_layers + + elif self.config.recompute_method == "block": + # Checkpoint the input activation of only a set number of individual + # Transformer layers and skip the rest. + # A method fully use the device memory removing redundant re-computation. + recompute_skip_num_layers = 0 + for layer_idx in range(self.num_layers_per_pipeline_rank): + # Skip recomputation when input grad computation is not needed. + # Need to have at least one input tensor with gradient computation + # for re-enterant autograd engine. + if self.config.fp8 and not hidden_states.requires_grad: + recompute_skip_num_layers += 1 + if ( + layer_idx >= recompute_skip_num_layers + and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers + ): + hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) + else: + hidden_states, context = custom(layer_idx, layer_idx + 1)( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + visual_pos_masks, + *deepstack_visual_embeds_tuple, + ) + else: + raise ValueError("Invalid activation recompute method.") + + return hidden_states + + def forward( + self, + hidden_states: Tensor | WrappedTensor, + attention_mask: Tensor | None, + context: Tensor | None = None, + context_mask: Tensor | None = None, + rotary_pos_emb: Tensor | None = None, + rotary_pos_cos: Tensor | None = None, + rotary_pos_sin: Tensor | None = None, + attention_bias: Tensor | None = None, + inference_context: BaseInferenceContext | None = None, + packed_seq_params: PackedSeqParams | None = None, + sequence_len_offset: Tensor | None = None, + *, + inference_params: BaseInferenceContext | None = None, + # args for deepstack + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + ): + """ + Perform the forward pass through the transformer block. + + This method handles the core computation of the transformer, including + self-attention, optional cross-attention, and feed-forward operations. + + Args: + hidden_states (Union[Tensor, WrappedTensor]): Input tensor of shape [s, b, h] + where s is the sequence length, b is the batch size, and h is the hidden size. + Can be passed as a WrappedTensor during inference to avoid an obsolete + reference in the calling function. + attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking + self-attention. + context (Tensor, optional): Context tensor for cross-attention. + context_mask (Tensor, optional): Mask for cross-attention context + rotary_pos_emb (Tensor, optional): Rotary positional embeddings. + attention_bias (Tensor): Bias tensor for Q * K.T of shape in shape broadcastable + to [b, num_head, sq, skv], e.g. [1, 1, sq, skv]. + Used as an alternative to apply attention mask for TE cuDNN attention. + inference_context (BaseInferenceContext, optional): Parameters for inference-time + optimizations. + packed_seq_params (PackedSeqParams, optional): Parameters for packed sequence + processing. + + Returns: + Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape + [s, b, h], and optionally the updated context tensor if cross-attention is used. + """ + if self.pre_process and deepstack_visual_embeds is not None: + assert len(deepstack_visual_embeds) <= len( + self.layers + ), "the deepstack_visual_embeds should on the first pp-stage" + + inference_context = deprecate_inference_params(inference_context, inference_params) + + # Delete the obsolete reference to the initial input tensor if necessary + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if not self.pre_process: + # See set_input_tensor() + hidden_states = self.input_tensor + + # Viewless tensor. + # - We only need to create a viewless tensor in the case of micro batch + # size (mbs) == 1, since in this case, 'hidden_states.transpose()' + # above creates a view tensor, and '.contiguous()' is a pass-through. + # For mbs >= 2, '.contiguous()' creates a new tensor, eliminating + # the need to make it viewless. + # + # However, we don't explicitly check mbs == 1 here because + # make_viewless_tensor() has negligible overhead when its input + # is already viewless. + # + # - For the 'else' case above, calling make_viewless_tensor() here is + # likely redundant, since p2p_communication.py (likely originator) + # already creates viewless tensors. That said, make_viewless_tensor() + # is called here to be future-proof and corner-case-proof. + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + + if self.config.sequence_parallel: + rng_context = tensor_parallel.get_cuda_rng_tracker().fork() + else: + rng_context = nullcontext() + + # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), + # otherwise do nothing extra at the outer level + # if we are using other fp8 recipes, then the context manager enter&exit are free + # we can wrap fp8_context within the for loop over layers, so that we can fine-grained + # control which layer will be fp8 or bf16 + use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed + use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed + outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext() + + with rng_context, outer_fp8_context: + # Forward pass. + if self.config.recompute_granularity == "full" and self.training: + hidden_states = self._checkpointed_forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + use_inner_fp8_context=use_inner_fp8_context, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + else: + for l_no, layer in enumerate(self.layers): + inner_fp8_context = ( + get_fp8_context(self.config, layer.layer_number - 1) + if use_inner_fp8_context + else nullcontext() + ) + with self.offload_context, inner_fp8_context: + hidden_states, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + attention_bias=attention_bias, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + + if self.pre_process and deepstack_visual_embeds is not None: + assert l_no == layer.layer_number - 1 + if l_no in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[l_no], + ) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + + if ( + torch.is_grad_enabled() + and self.config.cpu_offloading + and self.group_prefetch_offload_commit_async is not None + ): + hidden_states = self.group_prefetch_offload_commit_async(hidden_states) + + # Final layer norm. + if self.final_layernorm is not None: + hidden_states = self.final_layernorm(hidden_states) + # TENorm produces a "viewed" tensor. This will result in schedule.py's + # deallocate_output_tensor() throwing an error, so a viewless tensor is + # created to prevent this. + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + + # If this TransformerBlock is empty, input and output hidden states will be the same node + # on the computational graph and will lead to unexpected errors in pipeline schedules. + if not self.pre_process and len(self.layers) == 0 and not self.final_layernorm: + hidden_states = hidden_states.clone() + + return hidden_states + + def _deepstack_process( + self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor + ): + hidden_states = hidden_states.transpose(0, 1).contiguous() + local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds + hidden_states[visual_pos_masks, :] = local_this + hidden_states = hidden_states.transpose(0, 1).contiguous() + return hidden_states diff --git a/vime_plugins/models/qwen3_vl.py b/vime_plugins/models/qwen3_vl.py new file mode 100644 index 000000000..de5782c6f --- /dev/null +++ b/vime_plugins/models/qwen3_vl.py @@ -0,0 +1,201 @@ +"""Native Qwen3-VL: Megatron language model and a replicated HF vision tower.""" + +from __future__ import annotations + +import torch +from megatron.core import mpu, tensor_parallel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.transformer.module import MegatronModule +from transformers import AutoConfig + +from .qwen3_5_vl_utils import build_packed_mrope_position_ids +from .qwen3_omni_moe import Qwen3OmniMoeGPTModel +from .qwen3_omni_transformer import split_deepstack_embeddings + + +def _load_vision_model(hf_config, config): + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel + + device = ( + torch.device("cpu") if config.use_cpu_initialization else torch.device("cuda", torch.cuda.current_device()) + ) + with device: + vision_model = Qwen3VLVisionModel._from_config(hf_config.vision_config, attn_implementation="sdpa") + vision_model.to(dtype=config.params_dtype) + if config.recompute_granularity == "full": + vision_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + for parameter in vision_model.parameters(): + parameter.tensor_model_parallel = False + parameter.partition_dim = -1 + parameter.partition_stride = 1 + return vision_model + + +class Qwen3VLModel(MegatronModule): + def __init__(self, config, language_layer_spec, hf_config, args, *, pre_process, post_process, vp_stage): + super().__init__(config=config) + self.pre_process = pre_process + self.post_process = post_process + self.image_token_id = hf_config.image_token_id + self.video_token_id = hf_config.video_token_id + self.vision_start_token_id = hf_config.vision_start_token_id + self.spatial_merge_size = hf_config.vision_config.spatial_merge_size + + text_config = hf_config.text_config + rope = getattr(text_config, "rope_parameters", None) or text_config.rope_scaling + config.mrope_section = list(rope["mrope_section"]) + config.position_embedding_type = "mrope" + config.rotary_base = rope.get("rope_theta", getattr(text_config, "rope_theta", args.rotary_base)) + config.apply_rope_fusion = False + # The main Omni GPT class is also usable with a dense Qwen layer spec: + # its additions are interleaved MRoPE and checkpoint-aware DeepStack. + self.language_model = Qwen3OmniMoeGPTModel( + config=config, + transformer_layer_spec=language_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type="mrope", + rotary_percent=args.rotary_percent, + rotary_base=config.rotary_base, + scatter_embedding_sequence_parallel=False, + vp_stage=vp_stage, + ) + self.model = torch.nn.Module() + self.model.visual = _load_vision_model(hf_config, config) if pre_process else None + self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights + + @property + def decoder(self): + return self.language_model.decoder + + def shared_embedding_or_output_weight(self): + return self.language_model.shared_embedding_or_output_weight() + + def set_input_tensor(self, input_tensor): + self.language_model.set_input_tensor(input_tensor) + + def _inject_vision_embeddings(self, input_ids, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw): + embeddings = self.language_model.embedding(input_ids=input_ids, position_ids=None).transpose(0, 1).clone() + visual_mask = torch.zeros_like(input_ids, dtype=torch.bool) + positions, deepstack = [], [] + for values, grids, token_id in ( + (pixel_values, image_grid_thw, self.image_token_id), + (pixel_values_videos, video_grid_thw, self.video_token_id), + ): + mask = input_ids == token_id + if values is None: + if grids is not None or mask.any(): + raise ValueError("Qwen3-VL vision tokens/grids require matching pixel values") + continue + if grids is None: + raise ValueError("Qwen3-VL pixel values require matching grid_thw") + output = self.model.visual(values.to(dtype=self.model.visual.dtype), grid_thw=grids) + if hasattr(output, "pooler_output"): + features, layer_features = output.pooler_output, output.deepstack_features + else: + features, layer_features = output + if mask.sum().item() != features.shape[0]: + raise ValueError("Qwen3-VL token/features count mismatch") + embeddings[mask] = features.to(embeddings) + visual_mask |= mask + positions.append(mask.flatten().nonzero(as_tuple=False).flatten()) + deepstack.append(layer_features) + + deepstack_features = None + if deepstack: + # Images and videos are encoded separately but can interleave in a sample. + order = torch.cat(positions).argsort() + deepstack_features = [ + torch.cat(features)[order].to(embeddings) for features in zip(*deepstack, strict=True) + ] + embeddings = embeddings.transpose(0, 1).contiguous() + if self.config.sequence_parallel: + embeddings = tensor_parallel.scatter_to_sequence_parallel_region(embeddings).contiguous() + if deepstack_features is not None: + # Unlike Omni's frozen tower, this tower remains trainable. Each + # SP rank uses only part of DeepStack; sum its feature gradients + # before backpropagating through the replicated vision tower. + deepstack_features = [ + tensor_parallel.copy_to_tensor_model_parallel_region(features) for features in deepstack_features + ] + visual_mask, deepstack_features = split_deepstack_embeddings( + visual_mask, + deepstack_features, + tp_size=mpu.get_tensor_model_parallel_world_size(), + tp_rank=mpu.get_tensor_model_parallel_rank(), + sequence_parallel=True, + ) + return embeddings, visual_mask if deepstack_features is not None else None, deepstack_features + + def forward( + self, + input_ids, + position_ids=None, + attention_mask=None, + labels=None, + packed_seq_params=None, + loss_mask=None, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + **kwargs, + ): + if packed_seq_params is None or packed_seq_params.qkv_format != "thd": + raise ValueError("Qwen3-VL native training requires THD packed sequences") + if position_ids is None: + cu_seqlens = packed_seq_params.cu_seqlens_q_padded + if cu_seqlens is None: + cu_seqlens = packed_seq_params.cu_seqlens_q + position_ids = build_packed_mrope_position_ids( + input_ids, + cu_seqlens, + image_grid_thw, + video_grid_thw, + image_token_id=self.image_token_id, + video_token_id=self.video_token_id, + vision_start_token_id=self.vision_start_token_id, + spatial_merge_size=self.spatial_merge_size, + ) + embeddings, visual_mask, deepstack = self._inject_vision_embeddings( + input_ids, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw + ) + self.language_model.rotary_pos_emb.is_thd_format = True + return self.language_model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=embeddings, + labels=labels, + packed_seq_params=packed_seq_params, + loss_mask=loss_mask, + visual_pos_masks=visual_mask, + deepstack_visual_embeds=deepstack, + **kwargs, + ) + + +def get_qwen3_vl_model_provider(args, config, vp_stage): + """Use main's --spec provider interface without adding a Bridge branch.""" + if config.pipeline_model_parallel_size != 1 or config.context_parallel_size != 1: + raise ValueError("Qwen3-VL native training currently supports PP=1 and CP=1") + if args.mtp_num_layers: + raise ValueError("Qwen3-VL native MTP is not supported") + if args.transformer_impl != "transformer_engine": + raise ValueError("Qwen3-VL native training requires the TE/MindSpeed layer spec") + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + if hf_config.model_type != "qwen3_vl": + raise ValueError(f"{args.hf_checkpoint} is not a Qwen3-VL checkpoint") + layer_spec = get_gpt_layer_with_transformer_engine_spec(qk_layernorm=True) + + def model_provider(pre_process=True, post_process=True, vp_stage=None): + return Qwen3VLModel( + config, layer_spec, hf_config, args, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage + ) + + return model_provider diff --git a/vime_plugins/rollout_buffer/README.md b/vime_plugins/rollout_buffer/README.md index 1636a73ae..805aaed28 100644 --- a/vime_plugins/rollout_buffer/README.md +++ b/vime_plugins/rollout_buffer/README.md @@ -40,7 +40,7 @@ In addition, Rollout Buffer also provides some customizable functions to meet sp ### Example Script -First, you need to follow [Example: Qwen3-4B Model](../../docs/en/models/qwen3-4B.md) to configure the environment, download data and convert model checkpoints. And then run the following scripts: +First, you need to follow [Example: Qwen3-4B Model](../../docs/en/examples/qwen3-4B.md) to configure the environment, download data and convert model checkpoints. And then run the following scripts: ```bash cd vime_plugins/rollout_buffer bash rollout_buffer_example.sh diff --git a/vime_plugins/rollout_buffer/README_zh.md b/vime_plugins/rollout_buffer/README_zh.md index f17808a97..c1c4051cc 100644 --- a/vime_plugins/rollout_buffer/README_zh.md +++ b/vime_plugins/rollout_buffer/README_zh.md @@ -40,7 +40,7 @@ generator/ ### 示例脚本 -请仿照 [示例：Qwen3-4B 模型](../../docs/zh/models/qwen3-4B.md) 文档中配置好 vime 的运行环境，下载数据，并转换模型 ckpt。之后分别运行 +请仿照 [示例：Qwen3-4B 模型](../../docs/zh/examples/qwen3-4B.md) 文档中配置好 vime 的运行环境，下载数据，并转换模型 ckpt。之后分别运行 ```bash cd vime_plugins/rollout_buffer diff --git a/vime_plugins/rollout_buffer/rollout_buffer_example.py b/vime_plugins/rollout_buffer/rollout_buffer_example.py index 66167b957..aa73b214a 100644 --- a/vime_plugins/rollout_buffer/rollout_buffer_example.py +++ b/vime_plugins/rollout_buffer/rollout_buffer_example.py @@ -123,7 +123,7 @@ def log_raw_info(args, all_meta_info, rollout_id): wandb.log(log_dict) if args.use_tensorboard: - from vime.utils.tensorboard_utils import _TensorboardAdapter + from vime.observability.tensorboard_utils import _TensorboardAdapter tb = _TensorboardAdapter(args) tb.log(data=log_dict, step=step)