fix type annotation for ALL_ATTENTION_FUNCTIONS (#36690)

Corrects the type annotation to match actual usage. The variable was typed as
Dict[str, Dict[str, Callable]] but is actually used as Dict[str, Callable]
where keys are attention mechanism names and values are the corresponding
attention functions directly. This change makes the type annotation consistent
with how the dictionary is used in the codebase.

src/transformers/modeling_utils.py

@@ -5835,7 +5835,7 @@ def get_disk_only_shard_files(device_map, weight_map):
     return [fname for fname, devices in files_content.items() if set(devices) == {"disk"}]
 
 
-ALL_ATTENTION_FUNCTIONS: Dict[str, Dict[str, Callable]] = {}
+ALL_ATTENTION_FUNCTIONS: Dict[str, Callable] = {}
 
 ALL_ATTENTION_FUNCTIONS.update(
     {