Skip to content

Conversation

@wang2yn84
Copy link
Collaborator

@wang2yn84 wang2yn84 commented Oct 13, 2025

The current weight mapping lives in the model file. Consider we have more backends or more modes, the file size will grow much bigger. What's more important, the weight mapping should not sit in the model file directly, it's more of a extra feature we added to the model for RL.

This PR does the following:

  1. Separates the mapping function to separate file outside of model.py.
  2. Use mixin to inject the API to the model, so ppl can still call llama3_model.to_hf_mapping like before.
  3. Move lora config out of the mapping config, sitting directly in vllm config now. Mapping config is more for weight sync from actor to rollout model, but lora config is higly customizable and applies to rollout engine for initialization.
  4. Move MappingConfig from vllm_sampler to a separate file, together with mixin setup. Consider we have more rollout engines such as SGLang, they should be able to use the same MappingConfig and Mixin class.
  5. Provide the option for users to overwrite the mapping function from outside, e.g. MaxText.

Reference

Colab Notebook

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and all unit tests pass.
  • I have added all appropriate doc-strings/documentation.
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have signed the Contributor License Agreement.
  • I have followed Contribution Guidelines.

@wang2yn84 wang2yn84 mentioned this pull request Oct 13, 2025
6 tasks
@copybara-service copybara-service bot merged commit 4858530 into main Oct 13, 2025
7 checks passed
jxiong21029 pushed a commit to PLAN-Lab/tunix that referenced this pull request Nov 1, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants