[Raiden Weight Sync 4/7] Target-free and streaming weight conversion for vLLM

[YixuanWang-99]

Overview

Part of the stacked Raiden weight-sync enablement PR chain replacing #5089.

Stack:

• [M7] [Raiden Weight Sync 1/7] Add cross-repo drift guard and FFI resolution tests #5166: Cross-repo drift guard
• [M1] [Raiden Weight Sync 2/7] Support inhomogeneous layer cycle interval in raiden_unscan #5167: Inhomogeneous layer cycle support in raiden_unscan
• [M2] [Raiden Weight Sync 3/7] MoE 128-lane layout and padding for TPU GMM kernels #5168: MoE 128-lane layout and padding (pairs with Tunix T5)
• [M3] AI-Hypercomputer/maxtext (this PR): Target-free streaming weight conversion
• [M5] AI-Hypercomputer/maxtext: Shared expert gate weight_dtype fix
• [M4] AI-Hypercomputer/maxtext: TrainingEngine Raiden FFI integration (pairs with Tunix T2a)
• [M6] AI-Hypercomputer/maxtext: Rollout cleanup

Details

• Implements TargetFreeWeightConverter to convert MaxText weights directly into target vLLM formats without requiring pre-initialized target model structures.
• Implements streaming parameter conversion to minimize host peak RSS during staging.
• Adds comprehensive unit tests in tests/post_training/unit/weight_converter_test.py.

Verification

• pytest tests/post_training/unit/weight_converter_test.py (34/34 passed).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces target-free weight conversion and streaming weight conversion (convert_streaming) to the MaxText-to-MaxText converter, allowing incremental processing and reduced host memory usage. It also refactors compatibility context management and adds comprehensive unit tests. The review feedback correctly identifies several critical issues where assumptions about self.scan_axis being 1 are hardcoded, which could lead to incorrect shapes or runtime errors if configured differently. Additionally, the feedback points out a risky string prefix check for layer_ that could inadvertently match layers_ and cause crashes, suggesting a safer regex-based approach.

[codecov]

Codecov Report

❌ Patch coverage is 73.15789% with 51 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/integration/vllm/weight_converter.py	73.22%	36 Missing and 13 partials ⚠️
src/maxtext/integration/vllm/convert_utils.py	71.42%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!