diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 213b91e3a115..a6b9a517b20d 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -13,7 +13,6 @@ # limitations under the License. import torch import torch.nn as nn -import triton from torch.nn import functional as F from ..activations import ACT2FN @@ -159,6 +158,11 @@ def _load_deepgemm_kernel(): _deepgemm_available = True +def _cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return (a + b - 1) // b + + def w8a8_fp8_matmul( A: torch.Tensor, B: torch.Tensor, @@ -603,8 +607,8 @@ def __init__( if self.has_gate: gu_proj_out, gu_proj_in = 2 * self.intermediate_dim, self.hidden_dim self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, gu_proj_out, gu_proj_in, dtype=dtype)) - gu_scale_out = triton.cdiv(gu_proj_out, self.block_size[0]) if self.block_size is not None else 1 - gu_scale_in = triton.cdiv(gu_proj_in, self.block_size[1]) if self.block_size is not None else 1 + gu_scale_out = _cdiv(gu_proj_out, self.block_size[0]) if self.block_size is not None else 1 + gu_scale_in = _cdiv(gu_proj_in, self.block_size[1]) if self.block_size is not None else 1 self.gate_up_proj_scale_inv = nn.Parameter( torch.empty(self.num_experts, gu_scale_out, gu_scale_in, dtype=torch.float32) ) @@ -612,8 +616,8 @@ def __init__( else: u_proj_out, u_proj_in = self.intermediate_dim, self.hidden_dim self.up_proj = nn.Parameter(torch.empty(self.num_experts, u_proj_out, u_proj_in, dtype=dtype)) - u_scale_out = triton.cdiv(u_proj_out, self.block_size[0]) if self.block_size is not None else 1 - u_scale_in = triton.cdiv(u_proj_in, self.block_size[1]) if self.block_size is not None else 1 + u_scale_out = _cdiv(u_proj_out, self.block_size[0]) if self.block_size is not None else 1 + u_scale_in = _cdiv(u_proj_in, self.block_size[1]) if self.block_size is not None else 1 self.up_proj_scale_inv = nn.Parameter( torch.empty(self.num_experts, u_scale_out, u_scale_in, dtype=torch.float32) ) @@ -621,8 +625,8 @@ def __init__( d_proj_out, d_proj_in = self.hidden_dim, self.intermediate_dim self.down_proj = nn.Parameter(torch.empty(self.num_experts, d_proj_out, d_proj_in, dtype=dtype)) - d_scale_out = triton.cdiv(d_proj_out, self.block_size[0]) if self.block_size is not None else 1 - d_scale_in = triton.cdiv(d_proj_in, self.block_size[1]) if self.block_size is not None else 1 + d_scale_out = _cdiv(d_proj_out, self.block_size[0]) if self.block_size is not None else 1 + d_scale_in = _cdiv(d_proj_in, self.block_size[1]) if self.block_size is not None else 1 self.down_proj_scale_inv = nn.Parameter( torch.empty(self.num_experts, d_scale_out, d_scale_in, dtype=torch.float32) ) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index eb092019b678..d58c9a52fd33 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -66,10 +66,12 @@ ) from .integrations.deepspeed import _load_state_dict_into_zero3_model from .integrations.eager_paged import eager_paged_attention_forward +from .integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel +from .integrations.moe import ALL_EXPERTS_FUNCTIONS from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward from .integrations.sdpa_paged import sdpa_attention_paged_forward @@ -1969,11 +1971,14 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in def get_correct_experts_implementation(self, requested_experts: str | None) -> str: applicable_experts = "grouped_mm" if requested_experts is None else requested_experts - if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm"]: + base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys())) + valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns] + valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1] + valid_experts_str = ", ".join(valid_experts_str_list) + if applicable_experts not in base_experts_fns: message = ( f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are ' - '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"` ' - 'and `"experts_implementation=deepgemm"`.' + f"{valid_experts_str}." ) raise ValueError(message) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 6a27b6b5e0fb..fab48f9ddb8a 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -2823,6 +2823,20 @@ def test_error_wrong_attn_implementation(self): self.assertTrue('The only possible arguments are `attn_implementation="eager"' in str(cm.exception)) + def test_registered_experts_implementation_is_valid(self): + from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS + + def custom_experts_forward(*args, **kwargs): + pass + + experts_implementation = "custom_experts" + model = BaseModel(PreTrainedConfig()) + + with patch.dict(ALL_EXPERTS_FUNCTIONS._global_mapping, {}, clear=False): + ALL_EXPERTS_FUNCTIONS.register(experts_implementation, custom_experts_forward) + + self.assertEqual(model.get_correct_experts_implementation(experts_implementation), experts_implementation) + def test_not_available_flash(self): if is_flash_attn_2_available(): self.skipTest(reason="Please uninstall flash-attn package to run test_not_available_flash")