Skip to content

fix: filter out router so all-linear can be tuned for LoRA target modules #142

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
willmj merged 3 commits into from
Apr 30, 2025
Merged

fix: filter out router so all-linear can be tuned for LoRA target modules #142

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ba7e108
filter out router
willmj Apr 29, 2025
54a7979
fmt
willmj Apr 29, 2025
eb8ce1d
add importlib
willmj Apr 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion plugins/accelerated-moe/src/fms_acceleration_moe/utils/checkpoint_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,12 @@ def save_fsdp_optimizer(
(model_state_dict, optimizer_state_dict) = get_state_dict(model, optimizer)

# filter out lora state dict
# TODO: Once expert layers are supported for LoRA tuning
# remove the "router" filtering
lora_state_dict = {
k: v for k, v in model_state_dict.items() if "lora_A" in k or "lora_B" in k
k: v
for k, v in model_state_dict.items()
if ("lora_A" in k or "lora_B" in k) and "router" not in k
}

# - save model
Expand Down
1 change: 1 addition & 0 deletions plugins/accelerated-moe/tox.ini
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ envlist = py, lint
[testenv]
deps =
pytest>=7
importlib-metadata
Copy link
Collaborator

@anhuong anhuong Apr 30, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Noting this change: Error coming from linting failing to import from torch

.tox/py/lib/python3.9/site-packages/torch/distributed/elastic/rendezvous/__init__.py:142: in <module>
    from .registry import _register_default_handlers, _register_out_of_tree_handlers
.tox/py/lib/python3.9/site-packages/torch/distributed/elastic/rendezvous/registry.py:19: in <module>
    from importlib_metadata import entry_points
E   ModuleNotFoundError: No module named 'importlib_metadata'

Which looks like something that was added 7 months ago - https://github.com/pytorch/pytorch/blame/main/torch/distributed/elastic/rendezvous/registry.py#L18

So not sure why this is coming up as a new error when it didn't fail previously. But this should not affect the image build since we run on python 3.12 and this import only happens on python < 3.10.

-e {toxinidir}
skip_install = true
commands =
Expand Down