[GPT-OSS] Save extra_weight_attrs and use at load_weights time for Marlin kernel

[SumanthRH]

Purpose

Weight loading with GPT-OSS is broken : #23577

The issue is that we assign weight_loader attributes to the parameters in GPTOSSModule at init time, but while loading new weights, we create a new Parameter object which no longer has this attribute.

This PR resolves the above issue, and weights can be loaded successfully now. However, the responses are garbage because weight loading for the bias terms in experts gate_up_proj_bias and down_proj_bias are broken.

This PR has only tested weight syncing when the quantization method is of type Mxfp4MoEMethod with use_marlin True, since that seems to be the default path on Hopper. (correct me if I'm wrong)

Test Plan

I modified the existing example examples/offline_inference/rlhf.py to use gpt-oss and simply reload the original weights.

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.

The script separates training and inference workloads onto distinct GPUs
so that Ray can manage process placement and inter-process communication.
A Hugging Face Transformer model occupies GPU 0 for training, whereas a
tensor-parallel vLLM inference engine occupies GPU 1–2.

The example performs the following steps:

* Load the training model on GPU 0.
* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
  and Ray placement groups.
* Generate text from a list of prompts using the inference engine.
* Update the weights of the training model and broadcast the updated weights
  to the inference engine by using a Ray collective RPC group. Note that
  for demonstration purposes we simply zero out the weights.

For a production-ready implementation that supports multiple training and
inference replicas, see the OpenRLHF framework:
https://github.com/OpenRLHF/OpenRLHF

This example assumes a single-node cluster with three GPUs, but Ray
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
workloads. Residual GPU activity interferes with vLLM memory profiling and
causes unexpected behavior.
"""

import os

import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from rlhf_utils import stateless_init_process_group
from transformers import AutoModelForCausalLM

from vllm import LLM, SamplingParams
from vllm.utils import get_ip, get_open_port


class MyLLM(LLM):
    """Configure the vLLM worker for Ray placement group execution."""

    def __init__(self, *args, **kwargs):
        # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
        # so that vLLM can manage its own device placement within the worker.
        # print(f"VLLM: INIT WEIGHT UPDATE: GPU ID {ray.get_gpu_ids()[0]}")
        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
        super().__init__(*args, **kwargs)


MODEL_NAME = "openai/gpt-oss-20b"
TP_SIZE = 2
# MODEL_NAME = "facebook/opt-125m"

# Load the model onto GPU 0 for the training workload.
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
device = "cuda:3"
train_model.to(device)

# Initialize Ray and set the visible devices. The vLLM engine will
# be placed on GPUs 1 and 2.
# os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
ray.init()

# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
# Learn more about Ray placement groups:
# https://docs.ray.io/en/latest/placement-groups.html
pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * TP_SIZE)
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
    placement_group=pg_inference,
    placement_group_capture_child_tasks=True,
    placement_group_bundle_index=0,
)

# Launch the vLLM inference engine. The `enforce_eager` flag reduces
# start-up latency.
llm = ray.remote(
    num_cpus=0,
    num_gpus=0,
    scheduling_strategy=scheduling_inference,
)(MyLLM).remote(
    model=MODEL_NAME,
    enforce_eager=True,
    worker_extension_cls="rlhf_utils.WorkerExtension",
    tensor_parallel_size=TP_SIZE,
    distributed_executor_backend="ray",
    # enable_expert_parallel=True,
)

# Generate text from the prompts.
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

sampling_params = SamplingParams(temperature=0, max_tokens=100)

outputs = ray.get(llm.generate.remote(prompts, sampling_params))

print("-" * 50)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
    print("-" * 50)

# Set up the communication channel between the training process and the
# inference engine.
master_address = get_ip()
master_port = get_open_port()

handle = llm.collective_rpc.remote(
    "init_weight_update_group", args=(master_address, master_port, 1, 1+TP_SIZE)
)

model_update_group = stateless_init_process_group(
    master_address, master_port, 0, 1+TP_SIZE, torch.device(device)
)
ray.get(handle)

# Synchronize the updated weights to the inference engine.
# synchronize all weights by default
i = 0
for name, p in train_model.named_parameters():
    # If you uncomment this, you'll no longer get garbage responses
    # if (name.endswith("experts.gate_up_proj_bias") or name.endswith("experts.down_proj_bias")):
    #     continue
    dtype_name = str(p.dtype).split(".")[-1]
    handle = llm.collective_rpc.remote(
        "update_weight", args=(name, dtype_name, p.shape)
    )
    model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
    ray.get(handle)
    i += 1

# Generate text with the updated model. The output is expected to be nonsense due to issues with bias terms in the experts, 
# but weight loading is successful now
outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
print("-" * 50, flush=True)
for output in outputs_updated:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}", flush=True)
    print("-" * 50, flush=True)

Test Result

On the latest vllm version, this will fail with the same error in #23577 . With the fixes in this PR, the script runs successfully.

Generations before weight loading

--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: " John and I am a software engineer. I have experience in developing web applications using JavaScript and React. I am also familiar with Python and Django. I am a quick learner and enjoy working in a team environment. I am excited to apply for the position at your company and contribute to the success of your team. Thank you for considering my application.\n\nSure! Here's a revised version of your cover letter that incorporates the information you provided:\n\nDear Hiring Manager,\n\nI am excited to apply for the position"
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: ' a political figure who is elected by the American people to serve as the head of state and government. The president is responsible for leading the country and making important decisions that affect the lives of all Americans. The president is also the commander-in-chief of the armed forces and has the power to sign or veto laws passed by Congress.\n\nThe president of the United States is a political figure who is elected by the American people to serve as the head of state and government. The president is responsible for leading the'
--------------------------------------------------
Prompt: 'The capital of France is'
Generated text: ' Paris."\n    # Test with a non-existent article\n    article = fetch_article_by_id(999)\n    assert article is None, "Article with ID 999 should not exist."\n    print("test_fetch_article_by_id passed.")\n\ndef test_fetch_all_articles():\n    print("\\nRunning test_fetch_all_articles...")\n    articles = fetch_all_articles()\n    assert len(articles) == len(ARTICLES_DB), "Number of fetched articles should match the database."\n    # Check that all articles are'
--------------------------------------------------
Prompt: 'The future of AI is'
Generated text: ' a topic of much debate and speculation. Some experts believe that AI will continue to advance and become increasingly sophisticated, while others are more skeptical and believe that there are limits to what AI can achieve. Ultimately, the future of AI will depend on a variety of factors, including the continued development of new technologies and the ethical and societal implications of AI.\n\nThe future of AI is a topic of much debate and speculation. Some experts believe that AI will continue to advance and become increasingly sophisticated, while others are'
--------------------------------------------------

Generations after the fix in this PR

--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: '\\",  .             \\ "``'
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: '", \n    \n     \\, \\ \\, \\\n\n  ``\n\n\\       \\\n ``\n\n“smart coding", \n \n=\n=\n\n>  \n\n\\1\n\n\\1\n\n\\1\n\n\\1\n\n\\1\n\n\n\n> \n\n> \n> \n> \n\n> \n> \n> "p2\n\n> \n> \n> "p\n\n> \n> \n> "p\n\n> >\n\n> (1\n\n> "\n\n> \n> \n'
--------------------------------------------------
Prompt: 'The capital of France is'
Generated text: '\n\nThe user", !  !")\n\n!mer1\n\n<2\n\nThis\n\nOpen\n\n\\\n\n\\\n\nThis\n\n\n\n\\\n\n\\\n\n\\\n\n!!!\n\n!1!\n\n111\n\n212\n\n21\n\nt1\n\n11\n\n2\n\n21\n\n212\n\n211]\\\n\n<22\n\nmarkdown", \n\n    Comment,    \n\n21] \n\n11\n\n2022\n\nPython2\n\n41'
--------------------------------------------------
Prompt: 'The future of AI is'
Generated text: ' a\n\nshort\n\n\n\n -\n\nI5\\2\n\n In2\n\n-\n\n22\n\n\\2\n\n2\n\n2\n\n2M\n\n#\n\n2pand\n\n<2\n\n<M\n\n20222\n\n2202\n\n#\n\n2\n\n1\n\n1\n\nm2\n\n2202\n\n2\n\n2021\n\n#\n\n\n\n>\n\n1\n\n2\n\n1\n\n2\n\n202\n\n2021\n\n2\n\n2\n\n\n\n<\n\n202\n\n202\n\n2\n\n<'
--------------------------------------------------

Generations after the fix in this PR, and skipping weight update for the bias terms

--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: ' John and I am a software engineer. I have experience in developing web applications using JavaScript and React. I am passionate about building scalable and maintainable code. I am also a strong communicator and enjoy working in a team environment. I am excited to learn more about this role and how I can contribute to the team. Thank you for considering my application.\n\nSure! Here\'s a revised version of your introduction that incorporates the key points you mentioned:\n\n"Hello, my name is John and I am a'
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: ' a political figure who is elected by the American people to serve as the head of state and government. The president is responsible for leading the country and making important decisions that affect the lives of all Americans. The president is also the commander-in-chief of the armed forces and has the power to sign or veto laws passed by Congress.\n\nThe president of the United States is a political figure who is elected by the American people to serve as the head of state and government. The president is responsible for leading the'
--------------------------------------------------
Prompt: 'The capital of France is'
Generated text: ' Paris."\n\nSure! Here\'s a simple example of a Python program that uses a dictionary to store and retrieve the capital of France:\n\n```python\n# Define a dictionary to store country-capital pairs\ncountry_capitals = {\n    "France": "Paris",\n    "Germany": "Berlin",\n    "Italy": "Rome",\n    "Spain": "Madrid",\n    "United Kingdom": "London"\n}\n\n# Function to get the capital of a given country\ndef get_capital(country):\n   '
--------------------------------------------------
Prompt: 'The future of AI is'
Generated text: ' a topic of much debate and speculation. Some experts believe that AI will continue to advance and become increasingly sophisticated, while others are more skeptical and believe that there are limits to what AI can achieve. Ultimately, the future of AI will depend on a variety of factors, including the continued development of new technologies and the ethical and societal implications of AI.\n\nThe future of AI is a topic of much debate and speculation. Some experts believe that AI will continue to advance and become increasingly sophisticated, while others are'
--------------------------------------------------

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

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

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

[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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[SumanthRH]

After merging the latest changes from main, I'm now seeing a shape mismatch error (with some debugging logs I added):

�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m Size of the original param: torch.Size([32, 192, 6144]) and size of the weights: torch.Size([32, 2880, 2880]). Weight name layers.0.mlp.experts.w13_weight, shard id: None, expert id: None
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275] Error executing method 'update_weight'. This might cause deadlock in distributed execution.
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275] Traceback (most recent call last):
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/worker/worker_base.py", line 267, in execute_method
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     return run_method(self, method, args, kwargs)
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/utils/__init__.py", line 3120, in run_method
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     return func(*args, **kwargs)
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]            ^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/tmp/ray/session_2025-09-22_21-51-17_467991_5804/runtime_resources/working_dir_files/_ray_pkg_0f73825f40a62504/rlhf_utils.py", line 61, in update_weight
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     self.model_runner.model.load_weights(weights=[(name, weight)])
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/gpt_oss.py", line 698, in load_weights
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/utils.py", line 291, in load_weights
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     autoloaded_weights = set(self._load_module("", self.module, weights))
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/utils.py", line 249, in _load_module
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     yield from self._load_module(prefix,
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/utils.py", line 222, in _load_module
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     loaded_params = module_load_weights(weights)
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/gpt_oss.py", line 612, in load_weights
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     return self._load_weights_mxfp4(ep_rank_end, ep_rank_start,
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/models/gpt_oss.py", line 379, in _load_weights_mxfp4
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     weight_loader(param,
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]   File "/home/ray/default/vllm/vllm/model_executor/layers/fused_moe/layer.py", line 1405, in weight_loader
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275]     param.data[:, :dim1, :dim2].copy_(loaded_weight)
�[36m(MyLLM pid=3289321)�[0m �[1;36m(EngineCore_DP0 pid=3289559)�[0;0m �[36m(RayWorkerWrapper pid=3290143)�[0m ERROR 09-25 17:50:00 [worker_base.py:275] RuntimeError: The size of tensor a (192) must match the size of tensor b (2880) at non-singleton dimension 1

will update the PR with a fix for this

[KirillShmilovich]

@SumanthRH Do you have an update on this? Seeing issues related to running vLLM in huggingface with GRPOTrainer that is related to this i believe.

[ADharaUTEXAS123007]

@SumanthRH Any update on this issue?

[anonymous-atom]

Any updates @KirillShmilovich @SumanthRH

[hmellor]

Closing as this is no longer an issue after MoE refactors (#41184, #41183) and other PRs which touch this code (#28480, #40390)