`accelerate.load_checkpoint_and_dispatch` does not load GPT-OSS Models correctly

[KyleMylonakisProtopia]

System Info

- `Accelerate` version: 1.11.0
- Platform: Linux-6.8.0-87-generic-x86_64-with-glibc2.39
- `accelerate` bash location: /home/kyle/llama3/.venv/bin/accelerate
- Python version: 3.10.18
- Numpy version: 2.2.6
- PyTorch version: 2.8.0+cu128
- PyTorch accelerator: CUDA
- System RAM: 2015.50 GB
- GPU type: NVIDIA H100 80GB HBM3
- `Accelerate` default config:
        Not found

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

These two scripts differ in the weights they load

Run the following script with CUDA_VISIBLE_DEVICES=0 python <script_name>

import itertools

import torch
import transformers


def main() -> None:
    model = transformers.AutoModelForCausalLM.from_pretrained("/models/gpt-oss-20b", dtype=torch.bfloat16)

    uninitialized_tensors: list[str] = [
        name for name, tensor in itertools.chain(model.named_parameters(), model.named_buffers()) if (tensor == 0.0).all()
    ]
    names_str = ", ".join(uninitialized_tensors)
    print(f"The following tensors are all zeros and were likely not properly initialized: {names_str}")


if __name__ == "__main__":
    main()

Run the following script with CUDA_VISIBLE_DEVICES=0 torchrun --nproc_per_node=1 <script_name>

# © 2022-2025 Protopia AI, Inc. All rights reserved.

from __future__ import annotations

import itertools
import logging
import os
from typing import Final

import accelerate
import torch
import torch.distributed as dist
import transformers

_MODEL_PATH: Final[str] = "/models/gpt-oss-20b"
_TORCH_DTYPE: Final[torch.dtype] = torch.bfloat16
_ZERO_TENSOR_MSG: Final[str] = "The following tensors are all zeros and were likely not properly initialized: %s"


logger = logging.getLogger(__name__)


def _configure_logging() -> None:
    if logging.getLogger().hasHandlers():
        return
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )


def _init_distributed() -> int:
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    # torch.cuda.set_device(local_rank)
    return local_rank


def _collect_zero_tensors(model: torch.nn.Module) -> list[str]:
    zero_names: list[str] = []
    with torch.no_grad():
        for name, tensor in itertools.chain(model.named_parameters(), model.named_buffers()):
            if tensor.numel() == 0:
                continue
            if torch.count_nonzero(tensor).item() == 0:
                zero_names.append(name)
    return zero_names


def _gather_zero_tensors(local_zero_tensors: list[str]) -> list[str]:
    if not dist.is_initialized():
        return local_zero_tensors
    gathered: list[list[str]] = [None] * dist.get_world_size()  # type: ignore[list-item]
    dist.all_gather_object(gathered, local_zero_tensors)
    combined: list[str] = []
    for zero_list in gathered:
        combined.extend(zero_list)
    return combined


def main() -> None:
    """Check for uninitialized tensors across distributed ranks."""

    _configure_logging()
    local_rank = _init_distributed()
    device = torch.device(f"cuda:{local_rank}") if torch.cuda.is_available() else torch.device("cpu")

    logger.info("Rank %d loading model on %s", dist.get_rank(), device)

    with accelerate.init_empty_weights():
        model = transformers.AutoModelForCausalLM.from_pretrained(_MODEL_PATH, torch_dtype=_TORCH_DTYPE)
    model._apply(
            lambda t: torch.empty_like(t, device="cuda") if t.device == torch.device("meta") else t.to("cuda")
        )

    accelerate.load_checkpoint_and_dispatch(
        model,
        _MODEL_PATH,
        dtype=_TORCH_DTYPE,
        broadcast_from_rank0=True,
    )

    zero_tensors = _collect_zero_tensors(model)
    gathered_zero_tensors = sorted(set(_gather_zero_tensors(zero_tensors)))

    if dist.get_rank() == 0:
        if gathered_zero_tensors:
            logger.warning(_ZERO_TENSOR_MSG, ", ".join(gathered_zero_tensors))
        else:
            logger.info("All tensors are initialized across all ranks.")

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Expected behavior

These two scripts are loading a GPT-OSS model, however one is loading the model with accelerate. The model loaded with accelerate has the experts up and down projections initialized to zero.

This occurs because the weights saved in the state dict have different names form the modules in the Hugging Face transformes defined class due to the quantizatized form of the distributed model. The line of code which handles this is line 5424 in transformers/modeling_utils.py, which is not being called by accelerate.

[akshatvishu]

I think this is a known limitation of Accelerate when using the lower-level load_checkpoint_and_dispatch API directly on models that require complex weight name remapping (like GPT-OSS or quantized models).

Accelerate's load_checkpoint_in_model performs a direct mapping of the state dictionary. Crucially, if a key in the checkpoint does not match a key in the model definition exactly; it is treated as an "unexpected key" and skipped when strict=False (the default).

# src/accelerate/utils/modeling.py

if param_name not in model_keys:
    unexpected_keys.add(param_name)
    if not strict:
        continue  # Skip loading this parameter.

Because the keys are skipped the corresponding parameters in your model (which were initialized on the meta device via init_empty_weights) are never updated. They remain as empty/zero tensors, leading to the behavior you observed.

The easiest solution is to use transformer intergration:

model = transformers.AutoModelForCausalLM.from_pretrained(
    "/models/gpt-oss-20b",
    device_map="auto",  # Triggers Accelerate handles dispatch/offload automatically
    torch_dtype=torch.bfloat16
)

[KyleMylonakisProtopia]

The easiest solution is to use transformer intergration

Unfortunately, this isn't something we can really do in practice for our workflows because we use torch distributed modules to wrap our models for FSDP2 training.

Is there no plan to add support for known popular models like GPT-OSS? I've pinpointed the seeming exact quantization code that needs to be added in this ticket.

[akshatvishu]

Is there no plan to add support for known popular models like GPT-OSS? I've pinpointed the seeming exact quantization code that needs to be added in this ticket.

From what i can tell, accelerate’s policy is to stay format-agnostic, so the custom 4-bit unpack routine needed for GPT-OSS won’t be merged but we can wait for an official comment from the maintainers!

But I did some digging and found this script convert_gpt_oss_weights_to_hf.py in the transformer repo which you can use to do something like:
(Note: this is a concept and I didn't tested it myself!)

python -m transformers.models.gpt_oss.convert_gpt_oss_weights_to_hf \
    --input_dir  /path/to/raw/gpt-oss \
    --output_dir /path/to/gpt-oss-hf
# add --instruct if the model is instruct-tuned;
# Do NOT pass --mxfp4 if you want standard unpacked bfloat16 weights compatible with Accelerate

--instruct: If you are converting an instruct model, you must adds this flag. It ensures the eos_token_id is correctly set to 200002 (instead of the default 199999) and saves the correct Jinja2 chat template.

Then you can use Accelerate to load the standard weight natively :

# Use the converted path
_MODEL_PATH = "/path/to/converted/gpt-oss-hf" 

# ... init_empty_weights ...

# Accelerate can now load this because the keys match and weights are unpacked
accelerate.load_checkpoint_and_dispatch(
    model,
    _MODEL_PATH,
    dtype=torch.bfloat16,
    broadcast_from_rank0=True,
)

Also, keep in mind that, when using the converted model, ensure you load the tokenizer from the converted directory as well:

# Load from local converted path to get the correct EOS token (200002) and chat template
tokenizer = AutoTokenizer.from_pretrained("/path/to/converted/gpt-oss-hf")

Loading from the original Hub path may result in an incorrect eos_token_id for instruct models, causing generation issues.

[KyleMylonakisProtopia]

Thanks for digging; This could explain some discrepancies I was seeing by trying to manually convert the model weights by loading, and then using save_pretrained on the converted model.

[KyleMylonakisProtopia]

It seems like the file convert_gpt_oss_weights_to_hf.py is not actually distributed in the transformers releases, but only part of the repository.

@akshatvishu can you confirm this?

[akshatvishu]

Yeah! @KyleMylonakisProtopia Sorry that i forget to mentioned it earlier that to use it , easiest way is to build transformers from source! You can also wget the script only but then you'll have to take care of the dependencies manually!

I think it's intentionally not shipped in the wheel you get from pip install transformers. Conversion utilities live in the source tree (src/transformers/models/…) but are excluded from the PyPI package to keep the install small; they are one-off developer tools and not runtime code.

[KyleMylonakisProtopia]

The script doesn't seem to work out of the box for the current state of the GPT-OSS-20B model for either the root repo folder, or the original subfolder.

I'm going to see if these changes work, and if they do, I can see if I can open a PR in transformers.

[KyleMylonakisProtopia]

I opened a PR with Transformers:
huggingface/transformers#42901 and huggingface/transformers#42902