fix: add PIL image and processor serialization for VLM RPC#1070
fix: add PIL image and processor serialization for VLM RPC#1070Adiactive wants to merge 1 commit intoinclusionAI:mainfrom
Conversation
VLM training (e.g. geometry3k with Qwen2.5-VL) fails during RPC communication because PIL images and HuggingFace processors (ProcessorMixin) are not JSON-serializable. - Add SerializedPILImage: encodes PIL images as base64 PNG for rollout submit calls - Add SerializedProcessor: archives processors via save_pretrained into a zip, mirroring the existing SerializedTokenizer pattern - Wire both into serialize_value() and deserialize_value() - Add round-trip tests for both types Fixes: "Object of type JpegImageFile is not JSON serializable" Fixes: "Object of type Qwen2_5_VLProcessor is not JSON serializable"
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves critical RPC serialization failures encountered in multimodal training workloads. By introducing robust serialization mechanisms for PIL images and Hugging Face processors, it ensures that these complex data types can be correctly transmitted and stored within the system, thereby stabilizing VLM-related operations. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request effectively adds serialization support for PIL images and Hugging Face processors, which is crucial for VLM RPC workloads. The implementation for SerializedPILImage is clean and robust. The addition of SerializedProcessor correctly handles archiving and compression, and the integration into the main serialize_value and deserialize_value functions is well-executed. The accompanying tests are thorough and validate the new functionality well.
My main feedback is regarding significant code duplication between the new SerializedProcessor and the existing SerializedTokenizer classes. Refactoring this shared logic would greatly improve the maintainability of the code. I've left a specific comment with a suggestion on how to address this. I also have a minor suggestion to rename some constants for better clarity.
| raise ValueError(msg) | ||
|
|
||
|
|
||
| class SerializedProcessor(BaseModel): |
There was a problem hiding this comment.
There is significant code duplication between the new SerializedProcessor class and the existing SerializedTokenizer class. The methods for archiving (_archive_processor), compression (_maybe_compress), and decompression (_maybe_decompress) are nearly identical.
This duplication makes the code harder to maintain, as any bug fix or improvement in this logic would need to be applied in two places.
To improve maintainability, I recommend refactoring this common logic. For example, you could extract the shared logic into standalone utility functions that both SerializedTokenizer and SerializedProcessor can use.
| ) | ||
| compression = ( | ||
| zipfile.ZIP_STORED | ||
| if total_size < TOKENIZER_ARCHIVE_INLINE_THRESHOLD |
There was a problem hiding this comment.
Summary
This PR fixes VLM RPC serialization failures for multimodal training workloads by adding serialization support for:
ProcessorMixin) in checkpoint save metadataRoot cause
Two RPC payload types were not JSON-serializable in
areal/infra/rpc/serialization.py:JpegImageFile) in rollout submit pathQwen2_5_VLProcessor) inSaveLoadMetaduring save/checkpointChanges
SerializedPILImage(PNG bytes -> base64 -> round-trip to PIL image)SerializedProcessor(archive viasave_pretrained-> zip/base64 ->AutoProcessor.from_pretrained)serialize_value()anddeserialize_value()test_pil_image_roundtriptest_processorValidation
pytest tests/test_serialization.py -von GPU host: 10 passedruff check,ruff format --check)Training results after fix
Environment
Qwen2.5-VL-3B-InstructFollowing 1.0.2-vllm docker image
Script
Eval-Rollout/Reward
Related issue
Fixes #1069