[Core] Support weight_loader_v2 for UnquantizedLinearMethod

[kylesayrs]

Purpose

• Modernize UnquantizedLinearMethod by allowing the parameter to support weight_loader_v2
• This is useful for quantization methods which wrap unquantized parameters, such as transforms. See [Transform] [Quantization] Add transforms to compressed tensors #22486
  • These changes are not strictly required for CT transforms support, but will be required for supporting transforms + TP

Changes

• Change the weight Parameter to an instance of ModelWeightParameter
• Add UnquantizedLinearMethod to list of weight_loader_v2 supported methods

Caveats

• Unfortunately, torch.compile does not support capturing Parameter subclasses. I spend a couple hours trying to patch in support, and while it seems like there is theoretically support Tensor subclasses, Parameter subclasses are very difficult to support and patches lead to cascading errors.
  Typically, linear methods solve this by replacing the BasevLLMParameter with a torch.nn.Parameter after weight loading. However, this prevents the weight from being reloaded (see test_reload_weights_before_load_model).
  The simplest solution is to get a pointer to the weight data for execution and leave the original parameter untouched. Note that accessing data during execution leads to another torch.compile error.

# required by torch.compile
# do not overwrite with Parameter class to preserve weight reloading
self.layer_weight_data = layer.weight.data

• In most parameters, weight_loader_v2 is only an option for weight parameters, where as bias parameters always use weight_loader_v1. Unfortunately, for the QKVCrossParallelLinear class, both bias and weight parameters share the same weight loader function. Supporting weight_loader_v2 for bias parameters is out of scope of this PR, so the simplest solution is to implement a special bias_weight_loader for QKVCrossParallelLinear which is guaranteed to use weight_loader_v1.

def bias_weight_loader(self,
                       param: torch.nn.Parameter,
                       loaded_weight: torch.Tensor,
                       loaded_shard_id: Optional[str] = None):
    # just like all other parameters, does not yet
    # support loading bias with weight_loader_v2
    layer = (self.q_proj_decoder
             if loaded_shard_id == "q" else self.kv_proj_encoder)
    target_param = self.select_proj_params(layer, param)
    shard_id_args = (loaded_shard_id, ) if loaded_shard_id != "q" else ()
    layer.weight_loader(target_param, loaded_weight, *shard_id_args)

• Typically, a parameter's weight loader is passed from its parent module. However, some models override the weight loader. This is a bad pattern, as it is not compatible with quantization. Unfortunately, rectifying each of these models is out of scope for this PR, so the simplest solution is to support mutation of the weight_loader property with a note to begin supporting Model.load_weights.

def weight_loader(self) -> Callable:
    # NOTE(@ksayers) some models such as mamba_mixer2 override the
    # weight loader to support custom loading. In the future, model-specific
    # weight loading should be implemented via Model.load_weights. In the
    # meantime, support deleting and overriding `weight_loader`` attribute
    if self._weight_loader is None:
        raise AttributeError(f"{self.__class__.__name__} weight_loader "
                             "attribute has been deleted")

@weight_loader.setter
def weight_loader(self, value: Callable):
    self._weight_loader = value

@weight_loader.deleter
def weight_loader(self):
    self._weight_loader = None

Testing

• Tested meta-llama/Llama-3.1-8B-Instruct with TP ∈ {1, 2}
• Tested Qwen/Qwen1.5-MoE-A2.7B with TP ∈ {1, 2}
• CI testing

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request modernizes UnquantizedLinearMethod by adding support for weight_loader_v2. This is achieved by using ModelWeightParameter during weight creation and then converting it back to a standard torch.nn.Parameter after loading to ensure torch.compile compatibility. The changes are logical, well-contained, and correctly enable the new functionality. The implementation appears solid and aligns with the project's existing patterns.

[yewentao256]

Let's run the CI and see the results

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kylesayrs.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kylesayrs.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork