[Bug Fix] Fix naive_block_assignment always defaulting to False due to arg misalignment (#33848)
Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu>
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user