intel / intel/auto-round

Ensuring correct device dispatch after module replacement

Open
#1,308 3 comments 0 reactions 1 assignee Claimed by @yiliu30 View on GitHub
good first issue
Dominant language
Python
Stars
1.6k
Forks
175
Avg merge
1d 18h
Merged PRs (30d)
99

Description

When loading a model across multiple devices using device_map="auto", the accelerate library attaches hooks to move module inputs to the correct devices automatically.
For certain models (for example, gpt‑oss), we replace the original MoE modules during initialization. However, after this replacement, the newly inserted modules do not have the required Accelerate hooks attached. As a result, input tensors may not be dispatched to the correct devices.
We need to investigate when Accelerate binds these hooks during the model loading process. Based on that, we should re-dispatch (or re-attach hooks to) the replaced modules if necessary to ensure correct multi-device execution.

If needed, we can dispatch the replaced model with `dispatch_model`:
```python
from accelerate.big_modeling import dispatch_model
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.