Skip to content

Conversation

@abhinavclemson
Copy link
Collaborator

PiperOrigin-RevId: 825825400

Description

Refactor Tunix adapter to use a dedicated weight mapping utility.

This change introduces VllmWeightMapping in tunix/utils.py to encapsulate the logic for mapping MaxText weights to vLLM-compatible formats. The TunixMaxTextAdapter now uses an instance of VllmWeightMapping, which includes support for sharding knowledge and model-specific hook functions. New hook functions for Llama3.1 and Qwen3 are added to param_mapping.py

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.

@copybara-service copybara-service bot merged commit 1486d10 into main Nov 6, 2025
68 of 69 checks passed
@copybara-service copybara-service bot deleted the tunix-adaptor-refactor branch November 6, 2025 18:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants