Unverified Commit e1d97c38 authored by Runkai Tao's avatar Runkai Tao Committed by GitHub
Browse files

[Bug Fix] Fix `naive_block_assignment` always defaulting to False due to arg misalignment (#33848)


Signed-off-by: default avatarRunkai Tao <rt572@physics.rutgers.edu>
parent ec12d39d
......@@ -219,7 +219,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
self.max_loras,
self.adapter_enabled,
expert_map,
naive_block_assignment,
naive_block_assignment=naive_block_assignment,
)
moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora
......
......@@ -458,6 +458,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
adapter_enabled: torch.Tensor,
expert_map: torch.Tensor | None = None,
pad_sorted_ids: bool = False,
naive_block_assignment: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Aligns tokens and experts into block-sized chunks for LoRA-based
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment