diff --git a/docs/en/advanced/ascend-mxfp8-rollout.md b/docs/en/advanced/ascend-mxfp8-rollout.md new file mode 100644 index 000000000..a5c9621e3 --- /dev/null +++ b/docs/en/advanced/ascend-mxfp8-rollout.md @@ -0,0 +1,161 @@ +# Ascend W8A8-MXFP8 rollout + +Vime supports online W8A8-MXFP8 rollout weight updates on Ascend 950 through +vLLM and vLLM-Ascend. Training weights remain in BF16 during IPC or NCCL +transfer. Each rollout worker quantizes eligible Linear and fused MoE weights +before vLLM loads them, then vLLM-Ascend restores its inference layout. + +## Support matrix + +| Capability | Status | +| --- | --- | +| Ascend 950 W8A8-MXFP8 rollout | Supported | +| Colocated IPC online updates | Supported | +| Decoupled NCCL online updates | Supported | +| Linear and fused MoE weights described by ModelSlim | Supported | +| MXFP4 | Not supported in this feature | +| External vLLM server | Not supported in this feature | +| Disk/delta reload | Not supported in this feature | + +## Version baseline + +This integration targets the local `vllm-ascend/main` API at commit +`8bedc666`. It relies on: + +- `AscendModelSlimConfig`; +- the `W8A8_MXFP8` Linear and fused MoE schemes; +- `torch_npu.npu_dynamic_mx_quant`; +- vLLM's checkpoint-format weight-update transaction and worker-extension API. + +Use compatible vLLM, vLLM-Ascend, torch-npu, CANN, and ModelSlim versions for +your Ascend 950 environment. This Vime repository does not install the Ascend +software stack. + +## Model preparation + +The rollout checkpoint directory must contain the ModelSlim file +`quant_model_description.json`. Generate the complete file with the ModelSlim +version matched to your model and vLLM-Ascend environment. It must identify +the intended layers as `W8A8_MXFP8` and use group size 32. For example, a +complete model-specific file contains metadata and entries resembling: + +```json +{ + "quant_method": "ascend", + "group_size": 32, + "model.layers.0.self_attn.q_proj.weight": "W8A8_MXFP8", + "model.layers.0.self_attn.k_proj.weight": "W8A8_MXFP8", + "model.layers.0.input_layernorm.weight": "FLOAT" +} +``` + +The snippet is illustrative, not a complete description for a model. Every +required model layer must be represented using the names expected by the +current vLLM-Ascend ModelSlim parser. vLLM-Ascend rejects +`--quantization ascend` when the description file is missing or incompatible. + +The initial rollout checkpoint may contain MXFP8 inference weights, but the +online actor updates sent by Vime are BF16. Do not quantize these tensors on +the trainer side. + +## Configuration + +Enable the existing vLLM quantization argument: + +```bash +VLLM_ARGS=( + --vllm-quantization ascend + --vllm-gpu-memory-utilization 0.7 +) +``` + +### Colocated IPC + +Use the normal colocated topology and include `--colocate` in the Vime +invocation: + +```bash +python train.py \ + --colocate \ + --hf-checkpoint /path/to/rollout-checkpoint \ + --vllm-quantization ascend \ + --vllm-gpu-memory-utilization 0.7 \ + ... +``` + +Vime installs the MXFP8 worker extension and the native vLLM IPC transfer +backend. BF16 named tensors are shared through IPC and quantized in each +rollout worker. + +### Decoupled NCCL + +Use the normal non-colocated topology (omit `--colocate`): + +```bash +python train.py \ + --hf-checkpoint /path/to/rollout-checkpoint \ + --vllm-quantization ascend \ + --vllm-gpu-memory-utilization 0.7 \ + ... +``` + +Vime installs the same MXFP8 worker extension and selects the native vLLM NCCL +transfer backend. Only the transport changes; quantization and layout handling +remain worker-local and identical to the IPC path. + +## Weight-update lifecycle + +For both transports, Vime performs the following transaction: + +1. Start vLLM's checkpoint-format weight update, which restores model-format + metadata while preserving runtime tensor storage. +2. Prepare the Vime MXFP8 worker extension and wrap the common + `model.load_weights()` boundary. +3. Transfer BF16 tensors through IPC or NCCL. +4. Quantize eligible weights with `npu_dynamic_mx_quant` and load both the FP8 + weight and its `*_scale` tensor. +5. Finish native vLLM layerwise processing and reapply the vLLM-Ascend MXFP8 + inference layout. +6. Publish the new Vime weight version only after finish and finalization both + succeed. + +Other quantization configurations retain vLLM's existing behavior. + +## Ascend 950 validation + +The following acceptance procedure must be run in the target Ascend 950 +environment for both deployment modes: + +1. Start a small Vime recipe with `--update-weights-interval 1` and the MXFP8 + arguments above. +2. Confirm the vLLM-Ascend server loads the ModelSlim description and reaches + healthy state. +3. Allow at least two actor update/rollout iterations to complete. +4. Confirm each iteration completes the vLLM start/update/finish transaction + without missing-weight, missing-scale, shape, or dtype errors. +5. Confirm rollout generation succeeds after each update and that Vime reports + increasing weight versions. +6. Repeat once with `--colocate` (IPC) and once without it (NCCL). + +The CPU tests in this repository validate configuration, conversion, lifecycle, +and transport contracts with stubs. They do not replace this hardware test. + +## Troubleshooting + +`ModelSlim Quantization Config Not Found` +: Ensure `quant_model_description.json` is inside the directory passed through + `--hf-checkpoint` and was generated for the current model. + +Missing `weight_scale` or shape mismatch +: Verify all quantized layer names and fused-module mappings in the ModelSlim + file. Confirm `group_size` is 32 and the installed vLLM-Ascend matches the + baseline API. + +Worker extension conflict +: Remove a custom `--vllm-worker-extension-cls`. This feature requires Vime's + MXFP8 extension and fails closed instead of silently replacing another + extension. + +External server rejected +: Let Vime launch the rollout servers. Independently launched external vLLM + servers do not receive Vime's MXFP8 worker extension in this feature. diff --git a/docs/en/index.rst b/docs/en/index.rst index b6dbcab23..34d667c33 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -47,6 +47,8 @@ vime is built on `slime `_, the RL framework beh advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/low-precision.md + advanced/ascend-mxfp8-rollout.md advanced/vllm-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/docs/superpowers/plans/2026-07-14-ascend-mxfp8-rollout.md b/docs/superpowers/plans/2026-07-14-ascend-mxfp8-rollout.md new file mode 100644 index 000000000..279b2cbc4 --- /dev/null +++ b/docs/superpowers/plans/2026-07-14-ascend-mxfp8-rollout.md @@ -0,0 +1,438 @@ +# Ascend W8A8-MXFP8 Rollout Support Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add documented W8A8-MXFP8 online rollout weight updates for Ascend 950 to both Vime-managed IPC and NCCL vLLM-Ascend deployments. + +**Architecture:** A focused Ascend utility module lazily detects `AscendModelSlimConfig`, quantizes eligible BF16 Linear/MoE weights with `torch_npu.npu_dynamic_mx_quant`, and exposes unique worker-extension RPC methods. Vime's engine orchestration injects that extension when `--vllm-quantization ascend` is requested and composes it with vLLM's native checkpoint-format start/update/finish transaction, leaving IPC and NCCL wire formats unchanged. + +**Tech Stack:** Python 3.11+, PyTorch/torch-npu, vLLM weight transfer, vLLM-Ascend ModelSlim quantization, Ray RPC, pytest, Ruff, Sphinx/MyST documentation. + +## Global Constraints + +- Baseline vLLM-Ascend is local `main` commit `8bedc666`. +- Support only `W8A8_MXFP8`; MXFP4 is outside this plan. +- Support Vime-managed colocated IPC and decoupled NCCL; external servers and disk/delta reload are outside this plan. +- Keep `vllm_ascend` and `torch_npu` imports lazy so CPU, CUDA, and ROCm imports remain unaffected. +- Require the vLLM-Ascend ModelSlim `quant_model_description.json` contract with `group_size: 32`. +- Do not claim Ascend 950 hardware validation from this development machine. + +--- + +### Task 1: Ascend MXFP8 quantization utilities + +**Files:** +- Create: `vime/backends/vllm_utils/ascend_mxfp8.py` +- Create: `tests/utils/test_ascend_mxfp8.py` + +**Interfaces:** +- Produces: `is_ascend_mxfp8_config(quant_config: object) -> bool`. +- Produces: `quantize_mxfp8_weights(weights, model, dtype) -> Iterator[tuple[str, torch.Tensor]]`. +- Produces: `prepare_mxfp8_modules_for_reload(model) -> int` and `finalize_mxfp8_modules_after_reload(model) -> int`. +- Consumes: vLLM-Ascend schemes exposing `restore_weights_for_rl_loading(layer)` and `process_weights_after_loading(layer)`. + +- [ ] **Step 1: Write failing detection, mapping, quantization, and lifecycle tests** + +Create CPU-only tests using scoped `sys.modules` stubs. The core assertions must include: + +```python +def test_quantize_mxfp8_linear_weight_emits_weight_and_scale(fake_runtime): + model, source = fake_runtime.linear_model_and_weight() + output = list(mx.quantize_mxfp8_weights([("layers.0.proj.weight", source)], model, torch.bfloat16)) + assert [name for name, _ in output] == ["layers.0.proj.weight", "layers.0.proj.weight_scale"] + fake_runtime.dynamic_mx_quant.assert_called_once_with( + source.to(torch.bfloat16), axis=-1, dst_type=fake_runtime.float8_dtype + ) + + +def test_non_mxfp8_weight_passes_through_unchanged(fake_runtime): + model, source = fake_runtime.float_model_and_weight() + assert list(mx.quantize_mxfp8_weights([("layers.0.proj.weight", source)], model, torch.bfloat16)) == [ + ("layers.0.proj.weight", source) + ] + + +def test_prepare_uses_scheme_restore_when_native_metadata_has_not_restored_shapes(fake_runtime): + model, scheme = fake_runtime.transformed_model(native_shapes_restored=False) + assert mx.prepare_mxfp8_modules_for_reload(model) == 1 + scheme.restore_weights_for_rl_loading.assert_called_once() + + +def test_prepare_only_resets_marker_when_vllm_native_reload_already_restored_shapes(fake_runtime): + model, scheme = fake_runtime.transformed_model(native_shapes_restored=True) + assert mx.prepare_mxfp8_modules_for_reload(model) == 1 + scheme.restore_weights_for_rl_loading.assert_not_called() + assert model.layers[0].proj._mxfp8_transformed is False +``` + +Also cover `AscendModelSlimConfig` positive/negative detection, packed-module name mapping, fused MoE weights, scale flattening/squeezing, and idempotent finalize behavior. + +- [ ] **Step 2: Run the new test module and verify RED** + +Run: `pytest -q tests/utils/test_ascend_mxfp8.py` + +Expected: collection fails because `vime.backends.vllm_utils.ascend_mxfp8` does not exist. + +- [ ] **Step 3: Implement the minimal lazy utility module** + +Implement the following public structure, keeping optional imports inside functions: + +```python +MXFP8_QUANT_TYPE = "W8A8_MXFP8" + + +def is_ascend_mxfp8_config(quant_config: object) -> bool: + try: + from vllm_ascend.quantization.modelslim_config import AscendModelSlimConfig + except ImportError: + return False + return isinstance(quant_config, AscendModelSlimConfig) and MXFP8_QUANT_TYPE in getattr( + quant_config, "quant_description", {} + ).values() + + +def quantize_mxfp8_weights(weights, model, dtype=torch.bfloat16): + import torch_npu + + for name, value in weights: + if not _is_mxfp8_weight(name, model): + yield name, value + continue + quantized, scale = torch_npu.npu_dynamic_mx_quant( + value.to(dtype), axis=-1, dst_type=torch_npu.float8_e4m3fn + ) + scale = scale.flatten(-2, -1).squeeze(-1) + yield name, quantized + yield name + "_scale", scale +``` + +Implement `_module_from_param_name` using the model's `packed_modules_mapping`, recognize a scheme by the `restore_weights_for_rl_loading` capability, and make prepare/finalize handle both native-vLLM-restored shapes and still-transformed vLLM-Ascend tensors. + +- [ ] **Step 4: Run utility tests and verify GREEN** + +Run: `pytest -q tests/utils/test_ascend_mxfp8.py` + +Expected: all tests pass with no real `vllm_ascend` or `torch_npu` installation. + +- [ ] **Step 5: Run lint and commit Task 1** + +Run: `ruff check vime/backends/vllm_utils/ascend_mxfp8.py tests/utils/test_ascend_mxfp8.py` + +Expected: exit 0. + +Commit: + +```bash +git add vime/backends/vllm_utils/ascend_mxfp8.py tests/utils/test_ascend_mxfp8.py +git commit -m "feat: add Ascend MXFP8 rollout utilities" +``` + +### Task 2: vLLM worker extension lifecycle + +**Files:** +- Modify: `vime/backends/vllm_utils/ascend_mxfp8.py` +- Modify: `tests/utils/test_ascend_mxfp8.py` + +**Interfaces:** +- Consumes: the Task 1 utility functions. +- Produces: `AscendMXFP8WorkerExtension.prepare_ascend_mxfp8_weight_update() -> dict`. +- Produces: `AscendMXFP8WorkerExtension.finalize_ascend_mxfp8_weight_update(success: bool = True) -> dict`. + +- [ ] **Step 1: Write failing worker-extension tests** + +Add tests that build a fake worker with `model_runner.model.load_weights` and a fake MXFP8 config: + +```python +def test_worker_prepare_wraps_common_model_load_boundary(fake_worker): + result = fake_worker.prepare_ascend_mxfp8_weight_update() + fake_worker.model_runner.model.load_weights([("layers.0.proj.weight", torch.ones(2, 2))]) + assert result == {"active": True, "modules": 1} + assert fake_worker.original_loader_names == ["layers.0.proj.weight", "layers.0.proj.weight_scale"] + + +def test_worker_finalize_restores_original_loader_and_reapplies_layout(fake_worker): + original = fake_worker.model_runner.model.load_weights + fake_worker.prepare_ascend_mxfp8_weight_update() + result = fake_worker.finalize_ascend_mxfp8_weight_update(success=True) + assert fake_worker.model_runner.model.load_weights == original + assert result == {"active": True, "modules": 1} + + +def test_worker_finalize_failure_restores_loader_without_processing(fake_worker): + fake_worker.prepare_ascend_mxfp8_weight_update() + result = fake_worker.finalize_ascend_mxfp8_weight_update(success=False) + assert result == {"active": True, "modules": 0} + fake_worker.scheme.process_weights_after_loading.assert_not_called() +``` + +Also assert that non-MXFP8 ModelSlim configs return `{"active": False, "modules": 0}` and that nested prepare calls raise a clear `RuntimeError`. + +- [ ] **Step 2: Run worker-extension tests and verify RED** + +Run: `pytest -q tests/utils/test_ascend_mxfp8.py -k worker` + +Expected: failures because `AscendMXFP8WorkerExtension` is absent. + +- [ ] **Step 3: Implement the worker extension** + +Add a class with unique method names so vLLM's extension conflict guard accepts it: + +```python +class AscendMXFP8WorkerExtension: + def prepare_ascend_mxfp8_weight_update(self) -> dict: + model_runner = self.model_runner + if not is_ascend_mxfp8_config(model_runner.vllm_config.quant_config): + return {"active": False, "modules": 0} + if hasattr(self, "_vime_mxfp8_original_load_weights"): + raise RuntimeError("Ascend MXFP8 weight update is already active") + + model = model_runner.model + modules = prepare_mxfp8_modules_for_reload(model) + original_load_weights = model.load_weights + + def load_quantized(weights): + return original_load_weights( + quantize_mxfp8_weights(weights, model, model_runner.vllm_config.model_config.dtype) + ) + + self._vime_mxfp8_original_load_weights = original_load_weights + model.load_weights = load_quantized + return {"active": True, "modules": modules} + + def finalize_ascend_mxfp8_weight_update(self, success: bool = True) -> dict: + original = getattr(self, "_vime_mxfp8_original_load_weights", None) + if original is None: + return {"active": False, "modules": 0} + self.model_runner.model.load_weights = original + del self._vime_mxfp8_original_load_weights + modules = finalize_mxfp8_modules_after_reload(self.model_runner.model) if success else 0 + return {"active": True, "modules": modules} +``` + +- [ ] **Step 4: Run the whole utility/extension suite and verify GREEN** + +Run: `pytest -q tests/utils/test_ascend_mxfp8.py` + +Expected: all tests pass. + +- [ ] **Step 5: Commit Task 2** + +```bash +git add vime/backends/vllm_utils/ascend_mxfp8.py tests/utils/test_ascend_mxfp8.py +git commit -m "feat: add MXFP8 vLLM worker lifecycle" +``` + +### Task 3: Vime engine orchestration for IPC and NCCL + +**Files:** +- Modify: `vime/backends/vllm_utils/vllm_engine.py` +- Modify: `tests/utils/test_vllm_engine.py` + +**Interfaces:** +- Consumes: `vime.backends.vllm_utils.ascend_mxfp8.AscendMXFP8WorkerExtension` by qualified name. +- Produces: server args containing `worker_extension_cls` for both `weight_transfer_config.backend == "ipc"` and `"nccl"` when quantization is `ascend`. +- Produces: prepare/finalize collective RPC calls around native vLLM update transactions. + +- [ ] **Step 1: Write failing server-argument and transaction tests** + +Add parameterized tests: + +```python +@pytest.mark.parametrize(("colocate", "backend"), [(True, "ipc"), (False, "nccl")]) +def test_ascend_quantization_installs_same_worker_extension(vllm_args, monkeypatch, colocate, backend): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"quantization", "worker_extension_cls"})) + vllm_args.colocate = colocate + vllm_args.vllm_quantization = "ascend" + server_args, _ = mod._compute_server_args( + vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000 + ) + assert server_args["weight_transfer_config"]["backend"] == backend + assert server_args["worker_extension_cls"] == mod._ASCEND_MXFP8_WORKER_EXTENSION + + +def test_mxfp8_start_calls_native_then_prepare(vllm_engine, monkeypatch): + vllm_engine._uses_ascend_mxfp8 = True + calls = record_requests(vllm_engine, monkeypatch) + vllm_engine.start_weight_update(is_checkpoint_format=True) + assert calls == [ + ("start_weight_update", {"is_checkpoint_format": True}), + ("collective_rpc", {"method": "prepare_ascend_mxfp8_weight_update", "kwargs": {}}), + ] + + +def test_mxfp8_finish_promotes_pending_version_only_after_finalize(vllm_engine, monkeypatch): + vllm_engine._uses_ascend_mxfp8 = True + vllm_engine._weight_version = "old" + vllm_engine._pending_weight_version = "new" + calls = record_requests(vllm_engine, monkeypatch) + vllm_engine.finish_weight_update() + assert calls[-1] == ( + "collective_rpc", + {"method": "finalize_ascend_mxfp8_weight_update", "kwargs": {"success": True}}, + ) + assert vllm_engine._weight_version == "new" +``` + +Also test rejection of an independently launched external engine, conflict with a user-supplied worker extension, prepare/finalize failure cleanup, pending version behavior for both tensor and distributed update methods, and unchanged non-Ascend behavior. + +- [ ] **Step 2: Run focused engine tests and verify RED** + +Run: `pytest -q tests/utils/test_vllm_engine.py -k 'ascend or mxfp8'` + +Expected: failures for missing extension injection and lifecycle calls. + +- [ ] **Step 3: Implement extension selection and RPC orchestration** + +Add: + +```python +_ASCEND_MXFP8_WORKER_EXTENSION = ( + "vime.backends.vllm_utils.ascend_mxfp8.AscendMXFP8WorkerExtension" +) + + +def _requested_quantization(args, vllm_overrides: dict | None) -> str | None: + if vllm_overrides and "quantization" in vllm_overrides: + return vllm_overrides["quantization"] + return getattr(args, "vllm_quantization", None) +``` + +In `_compute_server_args`, after applying overrides, reject external rollout for `ascend`, reject a conflicting `worker_extension_cls`, and install `_ASCEND_MXFP8_WORKER_EXTENSION`. Ensure `_build_subprocess_env` includes the Vime package root whenever that extension is selected. + +In `VLLMEngine`, set `_uses_ascend_mxfp8`, track `_pending_weight_version`, and add a private collective-RPC helper. Native start must run before prepare. Native finish must run before successful finalize. On failures, call finalize with `success=False` without masking the original exception. Tensor/NCCL updates store a pending version under MXFP8 and retain immediate version updates for every existing backend. + +- [ ] **Step 4: Run engine tests and verify GREEN** + +Run: `pytest -q tests/utils/test_vllm_engine.py` + +Expected: all existing and new tests pass. + +- [ ] **Step 5: Run both transport contract suites** + +Run: + +```bash +pytest -q tests/utils/test_update_weight_from_tensor.py tests/utils/test_update_weight_from_distributed.py +``` + +Expected: all tests pass, proving both transports still use native start/update/finish wire formats. + +- [ ] **Step 6: Commit Task 3** + +```bash +git add vime/backends/vllm_utils/vllm_engine.py tests/utils/test_vllm_engine.py +git commit -m "feat: enable MXFP8 for IPC and NCCL rollout" +``` + +### Task 4: Feature support documentation + +**Files:** +- Create: `docs/en/advanced/ascend-mxfp8-rollout.md` +- Create: `docs/zh/advanced/ascend-mxfp8-rollout.md` +- Modify: `docs/en/index.rst` +- Modify: `docs/zh/index.rst` + +**Interfaces:** +- Documents: `--vllm-quantization ascend`, ModelSlim metadata, IPC/NCCL deployment modes, version baseline, limitations, troubleshooting, and hardware smoke validation. + +- [ ] **Step 1: Write the English and Chinese feature guides** + +Each guide must contain this support matrix and matching configuration example: + +```markdown +| Capability | Status | +| --- | --- | +| Ascend 950 W8A8-MXFP8 rollout | Supported | +| Colocated IPC online updates | Supported | +| Decoupled NCCL online updates | Supported | +| MXFP4 | Not supported in this feature | +| External vLLM server | Not supported in this feature | +| Disk/delta reload | Not supported in this feature | +``` + +```bash +VLLM_ARGS=( + --vllm-quantization ascend + --vllm-gpu-memory-utilization 0.7 +) +``` + +Document the requirement for `quant_model_description.json`, `W8A8_MXFP8`, and `group_size: 32`; list vLLM-Ascend baseline `8bedc666`; explain that BF16 crosses IPC/NCCL and online quantization occurs in each rollout worker. Include two-update smoke procedures for colocated and decoupled deployments and label them as commands to run on Ascend hardware, not locally verified results. + +- [ ] **Step 2: Add both guides to the Advanced Features toctrees** + +Add `advanced/ascend-mxfp8-rollout.md` to both `docs/en/index.rst` and `docs/zh/index.rst` immediately after `advanced/low-precision.md`; add `advanced/low-precision.md` to either index first if it is currently omitted. + +- [ ] **Step 3: Verify documentation structure** + +Run: + +```bash +python -m compileall -q vime +git diff --check +``` + +Expected: both commands exit 0. + +If the documentation dependencies are installed, also run `make -C docs html`; otherwise record the missing dependency and do not install unrelated packages. + +- [ ] **Step 4: Commit Task 4** + +```bash +git add docs/en/advanced/ascend-mxfp8-rollout.md docs/zh/advanced/ascend-mxfp8-rollout.md docs/en/index.rst docs/zh/index.rst +git commit -m "docs: add Ascend MXFP8 rollout guide" +``` + +### Task 5: Regression verification and delivery audit + +**Files:** +- Modify only if a verification failure reveals an in-scope defect. + +**Interfaces:** +- Verifies all preceding tasks and the approved design specification. + +- [ ] **Step 1: Run focused feature tests** + +Run: + +```bash +pytest -q \ + tests/utils/test_ascend_mxfp8.py \ + tests/utils/test_vllm_engine.py \ + tests/utils/test_update_weight_from_tensor.py \ + tests/utils/test_update_weight_from_distributed.py +``` + +Expected: zero failures. + +- [ ] **Step 2: Run static verification** + +Run: + +```bash +ruff check \ + vime/backends/vllm_utils/ascend_mxfp8.py \ + vime/backends/vllm_utils/vllm_engine.py \ + tests/utils/test_ascend_mxfp8.py \ + tests/utils/test_vllm_engine.py +python -m compileall -q vime +git diff --check +``` + +Expected: every command exits 0. + +- [ ] **Step 3: Audit scope and repository state** + +Run: + +```bash +git status --short --branch +git log --oneline --decorate main..HEAD +git diff --stat main...HEAD +``` + +Expected: branch is `rollout_support_mx`; committed changes are limited to the design/plan, MXFP8 utility and engine integration, tests, and bilingual documentation. Persistent `.planning/` working-memory files may remain untracked during execution but are not product deliverables. + +- [ ] **Step 4: Record the hardware-validation boundary** + +State in the final handoff that CPU mock and contract tests passed locally, while the documented Ascend 950 smoke procedure still needs execution in the target hardware environment. diff --git a/docs/superpowers/specs/2026-07-14-ascend-mxfp8-rollout-design.md b/docs/superpowers/specs/2026-07-14-ascend-mxfp8-rollout-design.md new file mode 100644 index 000000000..7600b7b9d --- /dev/null +++ b/docs/superpowers/specs/2026-07-14-ascend-mxfp8-rollout-design.md @@ -0,0 +1,133 @@ +# Ascend W8A8-MXFP8 Rollout Support Design + +## Goal + +Add online W8A8-MXFP8 rollout weight updates for Ascend 950 to Vime, using +the APIs and weight layouts provided by the local `vllm-ascend/main` +baseline. Both Vime-managed colocated CUDA/NPU IPC and decoupled NCCL weight +transfer must share the same quantization implementation. + +## Scope + +The feature supports: + +- Vime-managed vLLM rollout servers using vLLM-Ascend. +- `AscendModelSlimConfig` configurations containing `W8A8_MXFP8` layers. +- Online conversion of BF16 training weights with + `torch_npu.npu_dynamic_mx_quant`. +- Linear and fused MoE weights supported by vLLM-Ascend's W8A8-MXFP8 + schemes. +- Colocated IPC and decoupled NCCL online weight synchronization. + +The feature does not support MXFP4, disk/delta checkpoint reload, or an +independently launched external vLLM server in this iteration. + +## User Interface and Prerequisites + +Users enable the backend through Vime's existing vLLM argument forwarding: + +```bash +--vllm-quantization ascend +``` + +The rollout checkpoint directory must satisfy the local vLLM-Ascend main +contract. In particular, it must contain `quant_model_description.json` with +the intended layers marked `W8A8_MXFP8` and `group_size` set to 32. The model +weights supplied by the trainer remain BF16; Vime quantizes each online update +inside the rollout worker. + +## Architecture + +Vime will add an Ascend-specific MXFP8 worker adapter alongside its vLLM +backend. Imports of `vllm_ascend` and `torch_npu` will remain lazy so that the +existing CUDA and ROCm paths can import and run without those packages. + +Activation is fail-closed. The adapter is active only when the worker's +quantization configuration is an `AscendModelSlimConfig` and its quantization +description contains `W8A8_MXFP8`. Other Ascend quantization schemes and all +non-Ascend configurations retain the native vLLM behavior. + +The adapter hooks the worker-side weight loading lifecycle rather than either +transport implementation. Consequently, IPC and NCCL continue to carry BF16 +named tensors and both reach one MXFP8 implementation at the worker's +`model.load_weights()` boundary. + +## Weight-Update Lifecycle + +For each online update: + +1. Vime starts the native vLLM weight-update transaction. +2. The MXFP8 adapter walks supported vLLM-Ascend linear and fused MoE modules + and invokes `restore_weights_for_rl_loading()` where the inference layout + is currently transformed. +3. IPC or NCCL transfers BF16 named tensors without protocol changes. +4. At the common model-loading boundary, the adapter identifies parameters + whose target modules use the vLLM-Ascend W8A8-MXFP8 scheme. +5. Each eligible tensor is converted with + `torch_npu.npu_dynamic_mx_quant(..., axis=-1, + dst_type=torch_npu.float8_e4m3fn)`. The generated weight is loaded under + the original parameter name and its flattened uint8 scale is loaded using + the vLLM-Ascend `*_scale` parameter name. +6. BF16-only parameters pass through unchanged. +7. After every transfer chunk succeeds, the worker calls the applicable + vLLM-Ascend `process_weights_after_loading()` methods to recreate the + Ascend 950 inference layouts. +8. Only a successful native finish operation allows Vime to advance its + weight version. + +## Components + +The implementation will keep responsibilities separated: + +- An Ascend MXFP8 utility module detects the quantization configuration, + maps named parameters to target modules, produces quantized weight/scale + pairs, and performs layout restore/reapply operations. +- A small worker lifecycle integration installs the utility at vLLM server + startup and delegates transport to native vLLM. +- Existing IPC and NCCL trainer-side senders remain format-agnostic. Their + tests gain contracts showing both routes use the same worker lifecycle. +- A feature guide documents setup, supported modes, configuration, version + baseline, limitations, and hardware validation. + +## Error Handling + +- Optional Ascend dependencies are imported only after MXFP8 activation. +- Enabling `quantization=ascend` without a valid ModelSlim description is + rejected with a message naming the missing or incompatible configuration. +- Unsupported target parameter shapes or missing scale parameters raise + before the update is marked complete. +- Quantization or model loading errors preserve the original exception. The + post-load layout transformation and Vime weight-version advancement do not + run after a failed update. +- Repeated restore and process operations rely on vLLM-Ascend's idempotence + markers (`_mxfp8_transformed` and `_mxfp8_original_shapes`). + +## Testing and Validation + +CPU-runnable unit tests will use stubs for optional NPU dependencies and cover: + +- Positive and negative MXFP8 configuration detection. +- Linear and fused MoE parameter selection. +- BF16 passthrough for unquantized parameters. +- `npu_dynamic_mx_quant` arguments and generated `*_scale` names/shapes. +- Restore-before-load and process-after-load ordering. +- Failure behavior and idempotent lifecycle calls. +- IPC and NCCL contracts reaching the same MXFP8 integration without changing + their wire formats. + +Relevant existing Vime weight-transfer and vLLM backend tests will be rerun. +Static checks will cover all changed Python and documentation files. + +The local development machine has no Ascend 950 device, so completion will not +claim hardware execution. The feature guide will provide an Ascend end-to-end +smoke procedure that verifies initial server load, at least two online weight +updates, rollout generation after each update, and worker logs showing the +restore/quantize/reapply lifecycle for both IPC and NCCL deployments. + +## Compatibility + +The implementation baseline is the locally checked-out +`vllm-ascend/main` at design time (`8bedc666`). The integration will prefer +capability checks over broad version branching. CUDA FP8, ROCm, unquantized +rollout, existing compressed-tensor handling, and non-MXFP8 Ascend schemes +must retain their current behavior. diff --git a/docs/zh/advanced/ascend-mxfp8-rollout.md b/docs/zh/advanced/ascend-mxfp8-rollout.md new file mode 100644 index 000000000..ac0542008 --- /dev/null +++ b/docs/zh/advanced/ascend-mxfp8-rollout.md @@ -0,0 +1,138 @@ +# 昇腾 W8A8-MXFP8 Rollout + +Vime 通过 vLLM 和 vLLM-Ascend 支持昇腾 950 上的 W8A8-MXFP8 rollout 在线权重更新。 +训练权重通过 IPC 或 NCCL 传输时保持 BF16；每个 rollout worker 在 vLLM 加载权重前， +在线量化符合条件的 Linear 和 FusedMoE 权重，随后由 vLLM-Ascend 恢复推理布局。 + +## 支持矩阵 + +| 能力 | 状态 | +| --- | --- | +| 昇腾 950 W8A8-MXFP8 rollout | 支持 | +| Colocated IPC 在线更新 | 支持 | +| Decoupled NCCL 在线更新 | 支持 | +| ModelSlim 描述的 Linear 和 FusedMoE 权重 | 支持 | +| MXFP4 | 本特性暂不支持 | +| 外部独立 vLLM 服务 | 本特性暂不支持 | +| 磁盘/delta 权重重载 | 本特性暂不支持 | + +## 版本基线 + +本适配以本地 `vllm-ascend/main` 提交 `8bedc666` 的 API 为基线，依赖： + +- `AscendModelSlimConfig`； +- Linear 和 FusedMoE 的 `W8A8_MXFP8` scheme； +- `torch_npu.npu_dynamic_mx_quant`； +- vLLM checkpoint-format 权重更新事务和 worker extension API。 + +昇腾 950 环境需要安装相互兼容的 vLLM、vLLM-Ascend、torch-npu、CANN 和 ModelSlim。 +Vime 仓库本身不安装昇腾软件栈。 + +## 模型准备 + +Rollout checkpoint 目录必须包含 ModelSlim 生成的 `quant_model_description.json`。 +请使用与模型及 vLLM-Ascend 环境匹配的 ModelSlim 版本生成完整文件。目标量化层必须标记为 +`W8A8_MXFP8`，group size 必须为 32。完整的模型专用文件会包含类似以下的元数据和条目： + +```json +{ + "quant_method": "ascend", + "group_size": 32, + "model.layers.0.self_attn.q_proj.weight": "W8A8_MXFP8", + "model.layers.0.self_attn.k_proj.weight": "W8A8_MXFP8", + "model.layers.0.input_layernorm.weight": "FLOAT" +} +``` + +以上片段仅用于说明，并不是任一模型的完整描述。必须按照当前 vLLM-Ascend ModelSlim +解析器要求的名称覆盖所有必要层。缺失或不兼容的描述文件会导致 vLLM-Ascend 拒绝 +`--quantization ascend`。 + +初始 rollout checkpoint 可以包含 MXFP8 推理权重，但 Vime 在线发送的 actor 更新仍为 BF16； +不要在 trainer 侧对这些更新权重做 MXFP8 量化。 + +## 配置方法 + +使用现有的 vLLM 量化参数启用该特性： + +```bash +VLLM_ARGS=( + --vllm-quantization ascend + --vllm-gpu-memory-utilization 0.7 +) +``` + +### Colocated IPC + +在 Vime 启动命令中加入 `--colocate`： + +```bash +python train.py \ + --colocate \ + --hf-checkpoint /path/to/rollout-checkpoint \ + --vllm-quantization ascend \ + --vllm-gpu-memory-utilization 0.7 \ + ... +``` + +Vime 会安装 MXFP8 worker extension，并选择 vLLM 原生 IPC 权重传输。BF16 命名张量通过 +IPC 共享，在各 rollout worker 内完成在线量化。 + +### Decoupled NCCL + +使用普通的非 colocated 拓扑，不传 `--colocate`： + +```bash +python train.py \ + --hf-checkpoint /path/to/rollout-checkpoint \ + --vllm-quantization ascend \ + --vllm-gpu-memory-utilization 0.7 \ + ... +``` + +Vime 会安装同一个 MXFP8 worker extension，并选择 vLLM 原生 NCCL 权重传输。两种模式仅 +传输方式不同；在线量化和布局处理都在 worker 内执行，逻辑完全复用。 + +## 权重更新生命周期 + +两种传输模式都执行以下事务： + +1. 启动 vLLM checkpoint-format 权重更新，恢复模型格式元数据并保留运行时张量存储。 +2. 准备 Vime MXFP8 worker extension，包装统一的 `model.load_weights()` 入口。 +3. 通过 IPC 或 NCCL 传输 BF16 张量。 +4. 使用 `npu_dynamic_mx_quant` 在线量化目标权重，同时加载 FP8 权重和对应的 + `*_scale` 张量。 +5. 完成 vLLM 原生 layerwise processing，并恢复 vLLM-Ascend MXFP8 推理布局。 +6. 仅在 finish 和 finalization 都成功后发布新的 Vime weight version。 + +其他量化配置保持现有 vLLM 行为。 + +## 昇腾 950 验收步骤 + +以下步骤必须分别在目标昇腾 950 环境的 IPC 和 NCCL 模式执行： + +1. 使用 `--update-weights-interval 1` 和上述 MXFP8 参数启动一个小型 Vime 任务。 +2. 确认 vLLM-Ascend 成功加载 ModelSlim 描述并进入健康状态。 +3. 至少完成两轮 actor 更新和 rollout。 +4. 确认每轮 start/update/finish 事务中没有缺失权重、缺失 scale、shape 或 dtype 错误。 +5. 确认每次更新后 rollout 生成成功，且 Vime 报告的 weight version 持续递增。 +6. 使用 `--colocate` 完成一次 IPC 验收，再去掉该参数完成一次 NCCL 验收。 + +本仓库的 CPU 测试通过 stub 验证配置、量化转换、生命周期和传输契约，不能替代上述硬件测试。 + +## 故障排查 + +`ModelSlim Quantization Config Not Found` +: 确认 `--hf-checkpoint` 指向的目录包含 `quant_model_description.json`，且该文件针对当前模型生成。 + +缺失 `weight_scale` 或 shape 不匹配 +: 检查 ModelSlim 文件中的所有量化层名和融合模块映射；确认 `group_size` 为 32，且安装的 + vLLM-Ascend 与本文基线 API 兼容。 + +Worker extension 冲突 +: 移除自定义 `--vllm-worker-extension-cls`。本特性必须使用 Vime 的 MXFP8 extension；Vime + 会直接报错，不会静默替换已有扩展。 + +外部服务被拒绝 +: 让 Vime 自行拉起 rollout server。本特性不会向独立启动的外部 vLLM 服务注入 MXFP8 + worker extension。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 70fc4479c..de9d53c20 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -47,6 +47,8 @@ vime 构建于 `slime `_ 之上，slime 正是 G advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/low-precision.md + advanced/ascend-mxfp8-rollout.md advanced/vllm-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/tests/utils/test_ascend_mxfp8.py b/tests/utils/test_ascend_mxfp8.py new file mode 100644 index 000000000..b4e7dff53 --- /dev/null +++ b/tests/utils/test_ascend_mxfp8.py @@ -0,0 +1,317 @@ +"""CPU unit tests for Ascend W8A8-MXFP8 rollout weight reload helpers.""" + +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vime.backends.vllm_utils import ascend_mxfp8 as mx + + +class _FakeAscendModelSlimConfig: + def __init__(self, quant_description): + self.quant_description = quant_description + + +class _FloatConfig: + quant_description = {"quant_method": "ascend", "layer.weight": "FLOAT"} + + +class _FakeScheme: + def __init__(self): + self.restore_weights_for_rl_loading = MagicMock() + self.process_weights_after_loading = MagicMock() + + +class _QuantWrapper: + def __init__(self, scheme): + self.quant_method = scheme + + +class _Linear(torch.nn.Module): + def __init__(self, *, transformed: bool = False, native_shapes_restored: bool = True): + super().__init__() + self.scheme = _FakeScheme() + self.quant_method = _QuantWrapper(self.scheme) + if transformed and not native_shapes_restored: + weight_shape = (4, 2) + scale_shape = (2, 2, 2) + else: + weight_shape = (2, 4) + scale_shape = (2, 4) + self.weight = torch.nn.Parameter(torch.zeros(weight_shape), requires_grad=False) + self.weight_scale = torch.nn.Parameter(torch.zeros(scale_shape), requires_grad=False) + self._mxfp8_original_shapes = {"weight": (2, 4), "weight_scale": (2, 4)} + self._mxfp8_transformed = transformed + + +class _FloatLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(2, 4), requires_grad=False) + + +class _Block(torch.nn.Module): + def __init__(self, proj): + super().__init__() + self.proj = proj + + +class _Model(torch.nn.Module): + def __init__(self, proj): + super().__init__() + self.layers = torch.nn.ModuleList([_Block(proj)]) + self.packed_modules_mapping = {} + + +class _LoadableModel(_Model): + def __init__(self, proj): + super().__init__(proj) + self.loaded = [] + + def load_weights(self, weights): + loaded = list(weights) + self.loaded.extend(loaded) + return {name for name, _ in loaded} + + +class _PackedBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = _Linear() + + +class _PackedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_PackedBlock()]) + self.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + + +class _FusedMoE(torch.nn.Module): + def __init__(self): + super().__init__() + self.scheme = _FakeScheme() + self.quant_method = _QuantWrapper(self.scheme) + self.w13_weight = torch.nn.Parameter(torch.zeros(2, 4, 4), requires_grad=False) + self.w2_weight = torch.nn.Parameter(torch.zeros(2, 4, 4), requires_grad=False) + + +class _MoEBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.mlp = _FusedMoE() + + +class _MoEModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_MoEBlock()]) + self.packed_modules_mapping = {} + + +@pytest.fixture +def fake_ascend_config_module(monkeypatch): + package = types.ModuleType("vllm_ascend") + package.__path__ = [] + quantization = types.ModuleType("vllm_ascend.quantization") + quantization.__path__ = [] + modelslim = types.ModuleType("vllm_ascend.quantization.modelslim_config") + modelslim.AscendModelSlimConfig = _FakeAscendModelSlimConfig + monkeypatch.setitem(sys.modules, "vllm_ascend", package) + monkeypatch.setitem(sys.modules, "vllm_ascend.quantization", quantization) + monkeypatch.setitem(sys.modules, "vllm_ascend.quantization.modelslim_config", modelslim) + return _FakeAscendModelSlimConfig + + +@pytest.fixture +def fake_torch_npu(monkeypatch): + module = types.ModuleType("torch_npu") + module.float8_e4m3fn = torch.float8_e4m3fn + module.npu_dynamic_mx_quant = MagicMock( + return_value=(torch.zeros(2, 4, dtype=torch.float8_e4m3fn), torch.arange(6, dtype=torch.uint8).view(2, 3, 1)) + ) + monkeypatch.setitem(sys.modules, "torch_npu", module) + return module + + +@pytest.mark.unit +def test_detects_only_modelslim_config_with_mxfp8(fake_ascend_config_module): + assert mx.is_ascend_mxfp8_config( + fake_ascend_config_module({"quant_method": "ascend", "layers.0.proj.weight": "W8A8_MXFP8"}) + ) + assert not mx.is_ascend_mxfp8_config(fake_ascend_config_module({"quant_method": "ascend", "x": "FLOAT"})) + assert not mx.is_ascend_mxfp8_config(_FloatConfig()) + + +@pytest.mark.unit +def test_quantize_mxfp8_linear_weight_emits_weight_and_scale(fake_torch_npu): + model = _Model(_Linear()) + source = torch.ones(2, 4) + + output = list(mx.quantize_mxfp8_weights([("layers.0.proj.weight", source)], model, torch.bfloat16)) + + assert [name for name, _ in output] == ["layers.0.proj.weight", "layers.0.proj.weight_scale"] + fake_torch_npu.npu_dynamic_mx_quant.assert_called_once() + call = fake_torch_npu.npu_dynamic_mx_quant.call_args + assert torch.equal(call.args[0], source.to(torch.bfloat16)) + assert call.kwargs == {"axis": -1, "dst_type": torch.float8_e4m3fn} + assert output[1][1].shape == (2, 3) + + +@pytest.mark.unit +def test_non_mxfp8_weight_passes_through_unchanged(fake_torch_npu): + model = _Model(_FloatLinear()) + source = torch.ones(2, 4) + + output = list(mx.quantize_mxfp8_weights([("layers.0.proj.weight", source)], model, torch.bfloat16)) + + assert len(output) == 1 + assert output[0][0] == "layers.0.proj.weight" + assert output[0][1] is source + fake_torch_npu.npu_dynamic_mx_quant.assert_not_called() + + +@pytest.mark.unit +def test_packed_linear_source_name_maps_to_quantized_target(fake_torch_npu): + output = list( + mx.quantize_mxfp8_weights( + [("layers.0.q_proj.weight", torch.ones(2, 4))], + _PackedModel(), + ) + ) + + assert [name for name, _ in output] == ["layers.0.q_proj.weight", "layers.0.q_proj.weight_scale"] + + +@pytest.mark.unit +def test_fused_moe_source_path_maps_to_quantized_target(fake_torch_npu): + output = list( + mx.quantize_mxfp8_weights( + [("layers.0.mlp.experts.0.gate_proj.weight", torch.ones(2, 4))], + _MoEModel(), + ) + ) + + assert [name for name, _ in output] == [ + "layers.0.mlp.experts.0.gate_proj.weight", + "layers.0.mlp.experts.0.gate_proj.weight_scale", + ] + + +@pytest.mark.unit +def test_prepare_uses_scheme_restore_when_native_metadata_has_not_restored_shapes(): + layer = _Linear(transformed=True, native_shapes_restored=False) + layer.scheme.restore_weights_for_rl_loading.side_effect = lambda module: setattr( + module, "_mxfp8_transformed", False + ) + model = _Model(layer) + + assert mx.prepare_mxfp8_modules_for_reload(model) == 1 + layer.scheme.restore_weights_for_rl_loading.assert_called_once_with(layer) + assert layer._mxfp8_transformed is False + + +@pytest.mark.unit +def test_prepare_only_resets_marker_when_native_reload_already_restored_shapes(): + layer = _Linear(transformed=True, native_shapes_restored=True) + model = _Model(layer) + + assert mx.prepare_mxfp8_modules_for_reload(model) == 1 + layer.scheme.restore_weights_for_rl_loading.assert_not_called() + assert layer._mxfp8_transformed is False + + +@pytest.mark.unit +def test_finalize_reapplies_mxfp8_layout_idempotently(): + layer = _Linear(transformed=False) + layer.scheme.process_weights_after_loading.side_effect = lambda module: setattr( + module, "_mxfp8_transformed", True + ) + model = _Model(layer) + + assert mx.finalize_mxfp8_modules_after_reload(model) == 1 + assert mx.finalize_mxfp8_modules_after_reload(model) == 0 + layer.scheme.process_weights_after_loading.assert_called_once_with(layer) + + +def _make_worker(config, *, transformed: bool = True): + worker = mx.AscendMXFP8WorkerExtension() + layer = _Linear(transformed=transformed, native_shapes_restored=True) + layer.scheme.process_weights_after_loading.side_effect = lambda module: setattr( + module, "_mxfp8_transformed", True + ) + model = _LoadableModel(layer) + worker.model_runner = SimpleNamespace( + model=model, + vllm_config=SimpleNamespace( + quant_config=config, + model_config=SimpleNamespace(dtype=torch.bfloat16), + ), + ) + return worker, model, layer + + +@pytest.mark.unit +def test_worker_prepare_wraps_common_model_load_boundary(fake_ascend_config_module, fake_torch_npu): + config = fake_ascend_config_module({"layer.weight": "W8A8_MXFP8"}) + worker, model, _ = _make_worker(config) + + result = worker.prepare_ascend_mxfp8_weight_update() + loaded = model.load_weights([("layers.0.proj.weight", torch.ones(2, 4))]) + + assert result == {"active": True, "modules": 1} + assert loaded == {"layers.0.proj.weight", "layers.0.proj.weight_scale"} + assert [name for name, _ in model.loaded] == ["layers.0.proj.weight", "layers.0.proj.weight_scale"] + + +@pytest.mark.unit +def test_worker_finalize_restores_loader_and_reapplies_layout(fake_ascend_config_module): + config = fake_ascend_config_module({"layer.weight": "W8A8_MXFP8"}) + worker, model, layer = _make_worker(config) + original_loader = model.load_weights + + worker.prepare_ascend_mxfp8_weight_update() + result = worker.finalize_ascend_mxfp8_weight_update(success=True) + + assert model.load_weights == original_loader + assert result == {"active": True, "modules": 1} + layer.scheme.process_weights_after_loading.assert_called_once_with(layer) + + +@pytest.mark.unit +def test_worker_failure_cleanup_restores_loader_without_processing(fake_ascend_config_module): + config = fake_ascend_config_module({"layer.weight": "W8A8_MXFP8"}) + worker, model, layer = _make_worker(config) + original_loader = model.load_weights + + worker.prepare_ascend_mxfp8_weight_update() + result = worker.finalize_ascend_mxfp8_weight_update(success=False) + + assert model.load_weights == original_loader + assert result == {"active": True, "modules": 0} + layer.scheme.process_weights_after_loading.assert_not_called() + + +@pytest.mark.unit +def test_worker_non_mxfp8_config_is_noop(fake_ascend_config_module): + worker, model, _ = _make_worker(fake_ascend_config_module({"layer.weight": "FLOAT"})) + original_loader = model.load_weights + + assert worker.prepare_ascend_mxfp8_weight_update() == {"active": False, "modules": 0} + assert worker.finalize_ascend_mxfp8_weight_update() == {"active": False, "modules": 0} + assert model.load_weights == original_loader + + +@pytest.mark.unit +def test_worker_rejects_nested_prepare(fake_ascend_config_module): + worker, _, _ = _make_worker(fake_ascend_config_module({"layer.weight": "W8A8_MXFP8"})) + worker.prepare_ascend_mxfp8_weight_update() + + with pytest.raises(RuntimeError, match="already active"): + worker.prepare_ascend_mxfp8_weight_update() diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py index fccb5b4ff..46ec29149 100644 --- a/tests/utils/test_vllm_engine.py +++ b/tests/utils/test_vllm_engine.py @@ -305,6 +305,43 @@ def test_compute_server_args_no_sleep_mode_from_colocate(vllm_args): assert not getattr(vllm_args, "vllm_enable_sleep_mode", False) +@pytest.mark.unit +@pytest.mark.parametrize(("colocate", "backend"), [(True, "ipc"), (False, "nccl")]) +def test_ascend_quantization_installs_same_worker_extension(vllm_args, monkeypatch, colocate, backend): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"quantization", "worker_extension_cls"})) + vllm_args.rollout_external = False + vllm_args.colocate = colocate + vllm_args.vllm_quantization = "ascend" + + server_args, _ = mod._compute_server_args( + vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000 + ) + + assert server_args["weight_transfer_config"]["backend"] == backend + assert server_args["worker_extension_cls"] == mod._ASCEND_MXFP8_WORKER_EXTENSION + + +@pytest.mark.unit +def test_ascend_quantization_rejects_external_server(vllm_args, monkeypatch): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"quantization", "worker_extension_cls"})) + vllm_args.rollout_external = True + vllm_args.vllm_quantization = "ascend" + + with pytest.raises(ValueError, match="external vLLM"): + mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + + +@pytest.mark.unit +def test_ascend_quantization_rejects_conflicting_worker_extension(vllm_args, monkeypatch): + monkeypatch.setattr(mod, "_VLLM_SERVER_FIELDS", frozenset({"quantization", "worker_extension_cls"})) + vllm_args.rollout_external = False + vllm_args.vllm_quantization = "ascend" + vllm_args.vllm_worker_extension_cls = "custom.WorkerExtension" + + with pytest.raises(ValueError, match="worker extension"): + mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000) + + @pytest.mark.unit def test_get_base_gpu_id_colocate(vllm_args): vllm_args.colocate = True @@ -331,6 +368,25 @@ def fake_post(endpoint: str, payload: dict): assert calls[0][1] == {"is_checkpoint_format": True} +@pytest.mark.unit +def test_mxfp8_start_calls_native_then_prepare(vllm_engine, monkeypatch): + calls: list[tuple[str, dict]] = [] + vllm_engine._uses_ascend_mxfp8 = True + monkeypatch.setattr( + vllm_engine, + "_make_request", + lambda endpoint, payload: calls.append((endpoint, payload)) or {"ok": True}, + ) + + result = vllm_engine.start_weight_update(is_checkpoint_format=True) + + assert result == {"ok": True} + assert calls == [ + ("start_weight_update", {"is_checkpoint_format": True}), + ("collective_rpc", {"method": "prepare_ascend_mxfp8_weight_update", "kwargs": {}}), + ] + + @pytest.mark.unit def test_finish_weight_update_posts_empty_body(vllm_engine, monkeypatch): calls: list[tuple] = [] @@ -347,6 +403,58 @@ def fake_post(endpoint: str, payload: dict): assert calls == [("finish_weight_update", {})] +@pytest.mark.unit +def test_mxfp8_finish_finalizes_then_promotes_pending_version(vllm_engine, monkeypatch): + calls: list[tuple[str, dict]] = [] + vllm_engine._uses_ascend_mxfp8 = True + vllm_engine._weight_version = "old" + vllm_engine._pending_weight_version = "new" + monkeypatch.setattr( + vllm_engine, + "_make_request", + lambda endpoint, payload: calls.append((endpoint, payload)) or {"ok": True}, + ) + + result = vllm_engine.finish_weight_update() + + assert result == {"ok": True} + assert calls == [ + ("finish_weight_update", {}), + ( + "collective_rpc", + {"method": "finalize_ascend_mxfp8_weight_update", "kwargs": {"success": True}}, + ), + ] + assert vllm_engine._weight_version == "new" + assert vllm_engine._pending_weight_version is None + + +@pytest.mark.unit +def test_mxfp8_finish_failure_cleans_up_without_promoting_version(vllm_engine, monkeypatch): + calls: list[tuple[str, dict]] = [] + vllm_engine._uses_ascend_mxfp8 = True + vllm_engine._weight_version = "old" + vllm_engine._pending_weight_version = "new" + + def fake_request(endpoint, payload): + calls.append((endpoint, payload)) + if endpoint == "finish_weight_update": + raise RuntimeError("finish failed") + return {"ok": True} + + monkeypatch.setattr(vllm_engine, "_make_request", fake_request) + + with pytest.raises(RuntimeError, match="finish failed"): + vllm_engine.finish_weight_update() + + assert calls[-1] == ( + "collective_rpc", + {"method": "finalize_ascend_mxfp8_weight_update", "kwargs": {"success": False}}, + ) + assert vllm_engine._weight_version == "old" + assert vllm_engine._pending_weight_version is None + + @pytest.mark.unit def test_update_weights_from_tensor_posts_ipc_payload_and_records_version(vllm_engine, monkeypatch): posted: list[tuple[str, dict]] = [] @@ -377,6 +485,18 @@ def fake_post(endpoint: str, payload: dict): assert vllm_engine._weight_version == "42" +@pytest.mark.unit +def test_mxfp8_tensor_update_defers_version_until_finish(vllm_engine, monkeypatch): + vllm_engine._uses_ascend_mxfp8 = True + vllm_engine._weight_version = "old" + monkeypatch.setattr(vllm_engine, "_make_request", lambda endpoint, payload: {"ok": True}) + + vllm_engine.update_weights_from_tensor(names=[], dtype_names=[], shapes=[], ipc_handles=[], weight_version="new") + + assert vllm_engine._weight_version == "old" + assert vllm_engine._pending_weight_version == "new" + + @pytest.mark.unit def test_update_weights_from_tensor_does_not_advance_version_on_failure(vllm_engine, monkeypatch): """POST failure must not advance _weight_version (else a retry would skip the resync).""" @@ -449,6 +569,20 @@ def fake_make_request(endpoint: str, payload: dict) -> dict: assert vllm_engine._weight_version == "7" +@pytest.mark.unit +def test_mxfp8_distributed_update_defers_version_until_finish(vllm_engine, monkeypatch): + vllm_engine._uses_ascend_mxfp8 = True + vllm_engine._weight_version = "old" + monkeypatch.setattr(vllm_engine, "_make_request", lambda endpoint, payload: {"ok": True}) + + vllm_engine.update_weights_from_distributed( + [], [], [], group_name="vime", weight_version="new", packed=True + ) + + assert vllm_engine._weight_version == "old" + assert vllm_engine._pending_weight_version == "new" + + @pytest.mark.unit def test_get_url_ipv6_host(vllm_engine): vllm_engine.server_host = "[2001:db8::1]" diff --git a/vime/backends/vllm_utils/ascend_mxfp8.py b/vime/backends/vllm_utils/ascend_mxfp8.py new file mode 100644 index 000000000..9c7ee598e --- /dev/null +++ b/vime/backends/vllm_utils/ascend_mxfp8.py @@ -0,0 +1,189 @@ +"""Ascend W8A8-MXFP8 helpers for online rollout weight updates. + +Optional Ascend dependencies stay lazily imported so importing Vime on CPU, +CUDA, or ROCm does not require ``vllm-ascend`` or ``torch-npu``. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator + +import torch + +MXFP8_QUANT_TYPE = "W8A8_MXFP8" + + +def is_ascend_mxfp8_config(quant_config: object) -> bool: + """Return whether *quant_config* is a ModelSlim config containing MXFP8.""" + try: + from vllm_ascend.quantization.modelslim_config import AscendModelSlimConfig + except ImportError: + return False + + if not isinstance(quant_config, AscendModelSlimConfig): + return False + quant_description = getattr(quant_config, "quant_description", {}) + return MXFP8_QUANT_TYPE in quant_description.values() + + +def _mxfp8_scheme(module: object) -> object | None: + quant_method = getattr(module, "quant_method", None) + scheme = getattr(quant_method, "quant_method", quant_method) + if callable(getattr(scheme, "restore_weights_for_rl_loading", None)) and callable( + getattr(scheme, "process_weights_after_loading", None) + ): + return scheme + return None + + +def _module_from_param_name(model: object, name: str) -> object | None: + module_path = name.split(".")[:-1] + if not module_path: + return None + + packed_mapping = getattr(model, "packed_modules_mapping", {}) + reversed_mapping = { + original_name: fused_name + for fused_name, original_names in packed_mapping.items() + for original_name in original_names + } + if module_path[-1] in reversed_mapping: + module_path[-1] = reversed_mapping[module_path[-1]] + + current = model + try: + for part in module_path: + # Fused MoE checkpoints contain deeper expert paths even though the + # target vLLM module is already the fused quantized module. + if _mxfp8_scheme(current) is not None and hasattr(current, "w13_weight"): + return current + if isinstance(current, (torch.nn.ModuleList, torch.nn.Sequential)): + current = current[int(part)] + else: + current = getattr(current, part) + except (AttributeError, IndexError, KeyError, TypeError, ValueError): + return None + return current + + +def _is_mxfp8_weight(name: str, model: object) -> bool: + if not name.endswith("weight"): + return False + module = _module_from_param_name(model, name) + return module is not None and _mxfp8_scheme(module) is not None + + +def quantize_mxfp8_weights( + weights: Iterable[tuple[str, torch.Tensor]], + model: object, + dtype: torch.dtype = torch.bfloat16, +) -> Iterator[tuple[str, torch.Tensor]]: + """Convert eligible high-precision weights to Ascend MXFP8 on demand.""" + import torch_npu + + for name, value in weights: + if not _is_mxfp8_weight(name, model): + yield name, value + continue + + quantized, scale = torch_npu.npu_dynamic_mx_quant( + value.to(dtype), + axis=-1, + dst_type=torch_npu.float8_e4m3fn, + ) + scale = scale.flatten(-2, -1).squeeze(-1) + yield name, quantized + yield name + "_scale", scale + + +def _has_original_shapes(module: object) -> bool: + original_shapes = getattr(module, "_mxfp8_original_shapes", None) + if not isinstance(original_shapes, dict) or not original_shapes: + return False + for name, expected_shape in original_shapes.items(): + tensor = getattr(module, name, None) + if tensor is None or tuple(tensor.shape) != tuple(expected_shape): + return False + return True + + +def prepare_mxfp8_modules_for_reload(model: torch.nn.Module) -> int: + """Put MXFP8 modules into model-format layout before loading weights. + + Current vLLM restores model-format tensors from recorded metadata during + ``start_weight_update``. In that case only the vLLM-Ascend idempotence + marker needs resetting. Older/non-native paths still use the scheme's + explicit restore operation. + """ + prepared = 0 + for module in model.modules(): + scheme = _mxfp8_scheme(module) + if scheme is None: + continue + if getattr(module, "_mxfp8_transformed", False): + if _has_original_shapes(module): + module._mxfp8_transformed = False + else: + scheme.restore_weights_for_rl_loading(module) + prepared += 1 + return prepared + + +def finalize_mxfp8_modules_after_reload(model: torch.nn.Module) -> int: + """Reapply vLLM-Ascend inference layouts after a successful reload.""" + finalized = 0 + for module in model.modules(): + scheme = _mxfp8_scheme(module) + if scheme is None or getattr(module, "_mxfp8_transformed", False): + continue + scheme.process_weights_after_loading(module) + finalized += 1 + return finalized + + +class AscendMXFP8WorkerExtension: + """vLLM worker extension for online Ascend MXFP8 weight conversion. + + Method names are intentionally unique: vLLM rejects extension classes + whose attributes conflict with methods already defined by its worker. + """ + + def prepare_ascend_mxfp8_weight_update(self) -> dict[str, int | bool]: + model_runner = self.model_runner + if not is_ascend_mxfp8_config(model_runner.vllm_config.quant_config): + return {"active": False, "modules": 0} + if hasattr(self, "_vime_mxfp8_original_load_weights"): + raise RuntimeError("Ascend MXFP8 weight update is already active") + + model = model_runner.model + modules = prepare_mxfp8_modules_for_reload(model) + original_load_weights = model.load_weights + + def load_quantized(weights): + dtype = model_runner.vllm_config.model_config.dtype + return original_load_weights(quantize_mxfp8_weights(weights, model, dtype)) + + self._vime_mxfp8_original_load_weights = original_load_weights + model.load_weights = load_quantized + return {"active": True, "modules": modules} + + def finalize_ascend_mxfp8_weight_update(self, success: bool = True) -> dict[str, int | bool]: + original_load_weights = getattr(self, "_vime_mxfp8_original_load_weights", None) + if original_load_weights is None: + return {"active": False, "modules": 0} + + model = self.model_runner.model + model.load_weights = original_load_weights + del self._vime_mxfp8_original_load_weights + modules = finalize_mxfp8_modules_after_reload(model) if success else 0 + return {"active": True, "modules": modules} + + +__all__ = [ + "AscendMXFP8WorkerExtension", + "MXFP8_QUANT_TYPE", + "finalize_mxfp8_modules_after_reload", + "is_ascend_mxfp8_config", + "prepare_mxfp8_modules_for_reload", + "quantize_mxfp8_weights", +] diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 79c53f775..aa9188ea2 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -20,6 +20,7 @@ logger = logging.getLogger(__name__) _VLLM_WAKE_TAGS = frozenset({"weights", "kv_cache"}) +_ASCEND_MXFP8_WORKER_EXTENSION = "vime.backends.vllm_utils.ascend_mxfp8.AscendMXFP8WorkerExtension" def get_base_gpu_id(args, rank): @@ -66,13 +67,17 @@ def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]: env.setdefault("VLLM_SERVER_DEV_MODE", "1") if getattr(args, "vllm_enable_deterministic_inference", False): env["VLLM_BATCH_INVARIANT"] = "1" - if getattr(args, "colocate", False): + needs_vime_pythonpath = getattr(args, "colocate", False) or ( + server_args_dict.get("worker_extension_cls") == _ASCEND_MXFP8_WORKER_EXTENSION + ) + if needs_vime_pythonpath: import vime vime_root = os.path.dirname(os.path.dirname(os.path.abspath(vime.__file__))) existing_pp = env.get("PYTHONPATH", "") 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])) + if getattr(args, "colocate", False): env.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") worker_type = server_args_dict.get("_worker_type", "regular") @@ -130,6 +135,9 @@ def __init__( self.vllm_overrides = vllm_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine self._weight_version: str | None = None + self._pending_weight_version: str | None = None + quantization = _requested_quantization(args, self.vllm_overrides) + self._uses_ascend_mxfp8 = not args.rollout_external and str(quantization).lower() == "ascend" def init( self, @@ -279,7 +287,7 @@ def update_weights_from_tensor( if flush_cache: self.flush_cache() result = self._make_request("update_weights", {"update_info": payload}) - self._weight_version = str(weight_version) + self._record_weight_version(weight_version) return result def flush_cache(self): @@ -330,6 +338,12 @@ def get_weight_version(self): def set_weight_version(self, new_version: str): self._weight_version = str(new_version) + def _record_weight_version(self, weight_version: str) -> None: + if self._uses_ascend_mxfp8: + self._pending_weight_version = str(weight_version) + else: + self._weight_version = str(weight_version) + def release_memory_occupation(self, level: int = 2): self.flush_cache() response = requests.post(f"http://{self.server_host}:{self.server_port}/sleep", params={"level": level}) @@ -355,10 +369,49 @@ def init_weight_transfer_engine(self, payload: dict) -> dict: return self._make_request("init_weight_transfer_engine", payload) def start_weight_update(self, is_checkpoint_format: bool = False) -> dict: - return self._make_request("start_weight_update", {"is_checkpoint_format": is_checkpoint_format}) + result = self._make_request("start_weight_update", {"is_checkpoint_format": is_checkpoint_format}) + if self._uses_ascend_mxfp8: + try: + self._collective_rpc("prepare_ascend_mxfp8_weight_update") + except Exception: + try: + self._collective_rpc("finalize_ascend_mxfp8_weight_update", success=False) + except Exception: + logger.exception("Failed to clean up Ascend MXFP8 prepare state") + try: + self._make_request("finish_weight_update", {}) + except Exception: + logger.exception("Failed to close vLLM weight update after Ascend MXFP8 prepare error") + raise + return result def finish_weight_update(self) -> dict: - return self._make_request("finish_weight_update", {}) + if not self._uses_ascend_mxfp8: + return self._make_request("finish_weight_update", {}) + + try: + result = self._make_request("finish_weight_update", {}) + except Exception: + try: + self._collective_rpc("finalize_ascend_mxfp8_weight_update", success=False) + except Exception: + logger.exception("Failed to clean up Ascend MXFP8 state after vLLM finish error") + self._pending_weight_version = None + raise + + try: + self._collective_rpc("finalize_ascend_mxfp8_weight_update", success=True) + except Exception: + self._pending_weight_version = None + raise + + if self._pending_weight_version is not None: + self._weight_version = self._pending_weight_version + self._pending_weight_version = None + return result + + def _collective_rpc(self, method: str, **kwargs): + return self._make_request("collective_rpc", {"method": method, "kwargs": kwargs}) def update_weights_from_disk(self, model_path: str, load_format: str | None = None): del load_format @@ -413,7 +466,7 @@ def update_weights_from_distributed( "packed": bool(packed), } result = self._make_request("update_weights", {"update_info": update_info}) - self._weight_version = str(weight_version) + self._record_weight_version(weight_version) return result def pause_generation(self): @@ -479,6 +532,14 @@ def _normalize_vllm_wake_tags(tags: list[str] | None) -> list[str] | None: return normalized or None +def _requested_quantization(args, vllm_overrides: dict | None) -> str | None: + if vllm_overrides: + for key, value in vllm_overrides.items(): + if key.replace("-", "_") == "quantization": + return value + return getattr(args, "vllm_quantization", None) + + def _resolve_parallel_sizes(args, *, gpus_per_engine: int) -> tuple[int, int, int]: pp = int(getattr(args, "vllm_pipeline_parallel_size", 1) or 1) dp = int(getattr(args, "vllm_dp_size", None) or getattr(args, "vllm_data_parallel_size", 1) or 1) @@ -621,6 +682,20 @@ def _compute_server_args( if "model_path" in {k.replace("-", "_") for k in vllm_overrides}: kwargs["model"] = str(vllm_overrides.get("model_path") or vllm_overrides.get("model-path")) + if str(kwargs.get("quantization", "")).lower() == "ascend": + if args.rollout_external: + raise ValueError( + "Ascend MXFP8 rollout requires a Vime-managed vLLM server; " + "external vLLM servers are not supported." + ) + worker_extension = kwargs.get("worker_extension_cls") + if worker_extension and worker_extension != _ASCEND_MXFP8_WORKER_EXTENSION: + raise ValueError( + "Ascend MXFP8 rollout requires Vime's worker extension, but a different " + f"worker extension was configured: {worker_extension}" + ) + kwargs["worker_extension_cls"] = _ASCEND_MXFP8_WORKER_EXTENSION + # 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