Repository navigation
feat(ascend): add MXFP8 rollout support - #448
PierceZhou wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive performance metrics tracking for weight updates across both distributed and tensor-based backends, updates the weight transfer backend configuration to support Ascend (NPU) specific backends (npu_ipc and hccl), and aligns the weight update lifecycle with the vLLM 0.27 API. Additionally, it adds an address guard for online MXFP8 reloads to verify that graph-captured tensor storage is preserved. The reviewer feedback highlights a critical bug in update_weight_from_tensor.py where calculating tensor sizes can exhaust a lazy generator before it is sent, as well as timing metric inaccuracies in both backends due to lazy evaluation occurring outside timed blocks. Finally, the reviewer advises against leaving commented-out code in the megatron bridge initialization without explanation.
| while True: | ||
| export_started = time.perf_counter() | ||
| try: | ||
| hf_named_tensors = next(iterator) | ||
| except StopIteration: | ||
| break | ||
| export_seconds += time.perf_counter() - export_started |
There was a problem hiding this comment.
If the iterator returned by get_hf_weight_chunks yields generators or lazy iterables, calling sum(...) on hf_named_tensors at line 190 will exhaust the generator. This leaves hf_named_tensors empty when passed to _send_hf_params at line 194, silently failing the weight update. Additionally, the time spent converting/exporting the weights is currently excluded from export_seconds because the lazy evaluation happens during sum(...) outside the timed block. Converting next(iterator) to a list inside the timed block resolves both the generator exhaustion bug and the timing inaccuracy.
| while True: | |
| export_started = time.perf_counter() | |
| try: | |
| hf_named_tensors = next(iterator) | |
| except StopIteration: | |
| break | |
| export_seconds += time.perf_counter() - export_started | |
| while True: | |
| export_started = time.perf_counter() | |
| try: | |
| hf_named_tensors = list(next(iterator)) | |
| except StopIteration: | |
| break | |
| export_seconds += time.perf_counter() - export_started |
| while True: | ||
| export_started = time.perf_counter() | ||
| try: | ||
| hf_named_tensors = next(iterator) | ||
| except StopIteration: | ||
| break | ||
| self._perf_export_seconds += time.perf_counter() - export_started | ||
| if self._is_pp_src_rank: | ||
| hf_named_tensors = list(hf_named_tensors) |
There was a problem hiding this comment.
If get_hf_weight_chunks yields generators, the actual weight conversion/export work is deferred until the generator is evaluated (e.g., via list(hf_named_tensors)). Currently, list(hf_named_tensors) is called on line 246, which is outside the timed block for self._perf_export_seconds. This causes the export time metric to be highly inaccurate (underreported). Moving the list(...) conversion inside the timed block (only for the PP source rank to preserve the optimization on other ranks) ensures accurate performance metrics.
while True:
export_started = time.perf_counter()
try:
hf_named_tensors = next(iterator)
except StopIteration:
break
if self._is_pp_src_rank:
hf_named_tensors = list(hf_named_tensors)
self._perf_export_seconds += time.perf_counter() - export_started| @@ -1 +1 @@ | |||
| import vime_plugins.megatron_bridge.glm4v_moe # noqa: F401 # register GLM-4.6V bridge | |||
| #import vime_plugins.megatron_bridge.glm4v_moe # noqa: F401 # register GLM-4.6V bridge | |||
There was a problem hiding this comment.
Commenting out code instead of removing it or providing an explanatory comment/TODO is a maintainability anti-pattern. If the GLM-4.6V bridge registration is no longer needed, please delete this line entirely. If it is temporarily disabled, please add a TODO or an explanatory comment explaining why it is commented out and when it should be re-enabled.
Signed-off-by: PierceZhou <1342578551@qq.com>
0a7801f to
8b2275d
Compare
Signed-off-by: PierceZhou <1342578551@qq.com>
d6aad1c to
cc633a1
Compare
Signed-off-by: PierceZhou <1342578551@qq.com>
1 participant
Motivation
Add MXFP8 rollout support for the Ascend backend.
This enables vime to support rollout with MXFP8 models and improves compatibility with MXFP8 weight update workflows.
Changes
Implementation Details
Test
Tested with Ascend NPU environment.
Checklist
Integration evaluation (Ascend NPU)
Related quantization implementation: vllm-ascend#17644. These BF16/MXFP8 measurements exercise the integrated vIME rollout and weight-update workflow with online quantization in vLLM-Ascend.
Five short Qwen3-8B/32B training benchmarks showed 6%-11% lower total wall time with MXFP8 rollout. In the GSM8K-500 comparison, pass@4 changed from 95.0% to 93.8%,
mis_klfrom 0.0256 to 0.0349, and normalized ESS from 0.985 to 0.954. The math benchmark showed a larger quality trade-off: flexible extraction accuracy fell by about 6% relative for both Qwen3-8B and Qwen3-32B, despite higher inference throughput.A 400-step validation-mode run exercised rollout -> train -> weight_update over 8 epochs. It used
perform_rl_step=False, so it demonstrates execution of the workflow rather than policy improvement. These integrated measurements complement the PR's unit tests and CI; they are not a standalone measurement of the metric instrumentation added in this diff.