Remove relayouting linear weights logic from hf-adapters load-common.py - #280
Open
vinithakv wants to merge 1 commit into
Open
Remove relayouting linear weights logic from hf-adapters load-common.py#280vinithakv wants to merge 1 commit into
vinithakv wants to merge 1 commit into
Conversation
Author
|
Tested with and without patch and no significant regression found after removing relayout function from hf-adapters |
1 task
The relayouting is handled in torch_spyre/model_utils.py Signed-off-by: Vinitha Vijayan <vinithav@linux.vnet.ibm.com>
vinithakv
force-pushed
the
remove-relayout-redundant-logic
branch
from
August 5, 2026 15:46
920bbfa to
5eb92b6
Compare
vinithakv
marked this pull request as ready for review
August 5, 2026 15:47
arielge
reviewed
Aug 5, 2026
Comment on lines
1144
to
+1150
| if torch.device(DEVICE).type != "spyre": | ||
| model.to(dtype=dtype) | ||
| return | ||
|
|
||
| # Prime torch-spyre autoload before importing torch_spyre._C or calling | ||
| # torch.empty(..., device_layout=...). Calls with the spyre-only | ||
| # device_layout kwarg fail kwarg validation before dispatch. | ||
| torch.empty(1, device=DEVICE) | ||
|
|
||
| from torch_spyre._C import SpyreTensorLayout # type: ignore[import-not-found] | ||
|
|
||
| skip_layout_ptrs = _embedding_param_ids(model) | ||
|
|
||
| def _alloc_on_spyre(t: torch.Tensor) -> torch.Tensor: | ||
| # The row-major [1, 0] dim_order describes a 2-D permutation, so it only | ||
| # applies to 2-D matmul weights. 1-D tensors (norms, biases) and any | ||
| # higher-rank weight (e.g. the 3-D/4-D Conv2d and position-embedding | ||
| # tables in a multimodal checkpoint's vision/audio towers) keep the | ||
| # default layout — forcing [1, 0] on them raises "Incompatible host_size | ||
| # and dim_order". Embedding tables are gather-only and also skipped. | ||
| if t.dim() == 2 and t.data_ptr() not in skip_layout_ptrs: | ||
| stl = SpyreTensorLayout(t.shape, t.stride(), dtype, [1, 0]) | ||
| else: | ||
| stl = None | ||
| new: torch.Tensor = torch.empty( # type: ignore[call-overload] | ||
| t.shape, | ||
| device=torch.device(DEVICE), | ||
| device_layout=stl, | ||
| dtype=dtype, | ||
| ) | ||
| new.copy_(t.to(dtype)) | ||
| return new | ||
|
|
||
| for name, param in list(model.named_parameters()): | ||
| new = _alloc_on_spyre(param.data) | ||
| module_path, _, attr = name.rpartition(".") | ||
| owner = model.get_submodule(module_path) if module_path else model | ||
| setattr(owner, attr, nn.Parameter(new, requires_grad=False)) | ||
| from torch_spyre.model_utils import load_model_to_spyre | ||
|
|
||
| for name, buf in list(model.named_buffers()): | ||
| new = _alloc_on_spyre(buf) | ||
| module_path, _, attr = name.rpartition(".") | ||
| owner = model.get_submodule(module_path) if module_path else model | ||
| persistent = attr not in owner._non_persistent_buffers_set | ||
| owner.register_buffer(attr, new, persistent=persistent) | ||
| load_model_to_spyre(model, dtype=dtype) |
Contributor
There was a problem hiding this comment.
IIUC we do not need to explicitly import torch-spyre and call the method, and doing model.to(dtype=dtype, device=DEVICE) (for both cpu and spyre) should do this anyway under the hood
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The relayouting is handled in torch_spyre/model_utils.py