Padding free bufix ljl#202
Conversation
Add `_call_with_supported_kwargs` utility to filter out unsupported keyword arguments when calling forward methods, preventing errors from incompatible function signatures. This fixes issues where `origin_forward` methods may not accept all passed kwargs.
…handling - Add `_call_with_supported_kwargs` and `_call_create_causal_mask` helpers to filter unsupported kwargs - Rename `cache_position` parameter to `q_length` in flash_attention_mask and sdpa_mask for clarity - Fix device detection in sdpa_mask when `q_length` is not a tensor - Ensure compatibility with models that don't accept `cache_position` in causal mask functions
…ill path In sequence parallel training, when newer Transformers versions pass q_length/q_offset instead of cache_position, the causal mask creation may still see the local shard length. This change restores the global query length for the no-cache prefill path while keeping cache/sliding paths with their upstream offsets. Also refactor GDN padding-free detection to use transformers version check instead of source inspection, supporting transformers >= 5.9.0.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces compatibility helpers and updates masking and forward logic to support varying transformers library versions, specifically handling changes in argument signatures such as cache_position. The review feedback primarily focuses on optimizing performance by caching inspect.signature results and moving signature checks out of hot-path function definitions to avoid significant overhead during the model's forward pass.
| def _call_with_supported_kwargs(fn, *args, **kwargs): | ||
| signature = inspect.signature(fn) | ||
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): | ||
| kwargs = {key: value for key, value in kwargs.items() if key in signature.parameters} | ||
| return fn(*args, **kwargs) |
There was a problem hiding this comment.
Calling inspect.signature is a relatively expensive operation. Since _call_with_supported_kwargs is used within the model's forward pass (e.g., during mask creation), calling it repeatedly can introduce significant overhead. Consider caching the signature of the function to improve performance.
| def _call_with_supported_kwargs(fn, *args, **kwargs): | |
| signature = inspect.signature(fn) | |
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): | |
| kwargs = {key: value for key, value in kwargs.items() if key in signature.parameters} | |
| return fn(*args, **kwargs) | |
| def _call_with_supported_kwargs(fn, *args, **kwargs): | |
| if not hasattr(_call_with_supported_kwargs, '_cache'): | |
| _call_with_supported_kwargs._cache = {} | |
| sig = _call_with_supported_kwargs._cache.get(fn) | |
| if sig is None: | |
| sig = _call_with_supported_kwargs._cache[fn] = inspect.signature(fn) | |
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in sig.parameters.values()): | |
| kwargs = {key: value for key, value in kwargs.items() if key in sig.parameters} | |
| return fn(*args, **kwargs) |
| def sdpa_mask(batch_size, q_length=None, kv_length=None, *args, **kwargs): | ||
| origin_sdpa = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'] | ||
| origin_uses_cache_position = 'cache_position' in inspect.signature(origin_sdpa).parameters |
There was a problem hiding this comment.
The signature inspection and the lookup of the original SDPA function should be performed once outside of the sdpa_mask definition. This avoids redundant computation during every forward pass of the model.
| def sdpa_mask(batch_size, q_length=None, kv_length=None, *args, **kwargs): | |
| origin_sdpa = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'] | |
| origin_uses_cache_position = 'cache_position' in inspect.signature(origin_sdpa).parameters | |
| origin_sdpa = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] | |
| origin_uses_cache_position = 'cache_position' in inspect.signature(origin_sdpa).parameters | |
| def sdpa_mask(batch_size, q_length=None, kv_length=None, *args, **kwargs): |
| return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs) | ||
|
|
||
| return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs) |
There was a problem hiding this comment.
The calls to origin_sdpa should use _call_with_supported_kwargs to ensure compatibility across different Transformers versions. This is particularly important for the device argument, which may not be accepted by older versions of the SDPA mask function.
| return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs) | |
| return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs) | |
| return _call_with_supported_kwargs(origin_sdpa, batch_size, cache_position, kv_length, *args, **kwargs) | |
| return _call_with_supported_kwargs(origin_sdpa, batch_size, q_length, kv_length, *args, device=device, **kwargs) |
| if 'cache_position' in inspect.signature(masking_utils.origin_create_causal_mask).parameters: | ||
| cache_position_or_past_key_values = torch.arange( | ||
| 0, | ||
| input_embeds.shape[1], | ||
| device=input_embeds.device, | ||
| ) |
There was a problem hiding this comment.
Similar to the optimization suggested for sdpa_mask, the signature of masking_utils.origin_create_causal_mask should be inspected once outside the create_causal_mask function to avoid overhead in the forward pass.
| def _call_with_supported_kwargs(fn, *args, **kwargs): | ||
| signature = inspect.signature(fn) | ||
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): | ||
| kwargs = {key: value for key, value in kwargs.items() if key in signature.parameters} | ||
| return fn(*args, **kwargs) |
There was a problem hiding this comment.
Caching the signature in _call_with_supported_kwargs is recommended here as well, as this function is called within the forward pass of the GatedDeltaNet layers.
| def _call_with_supported_kwargs(fn, *args, **kwargs): | |
| signature = inspect.signature(fn) | |
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): | |
| kwargs = {key: value for key, value in kwargs.items() if key in signature.parameters} | |
| return fn(*args, **kwargs) | |
| def _call_with_supported_kwargs(fn, *args, **kwargs): | |
| if not hasattr(_call_with_supported_kwargs, '_cache'): | |
| _call_with_supported_kwargs._cache = {} | |
| sig = _call_with_supported_kwargs._cache.get(fn) | |
| if sig is None: | |
| sig = _call_with_supported_kwargs._cache[fn] = inspect.signature(fn) | |
| if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in sig.parameters.values()): | |
| kwargs = {key: value for key, value in kwargs.items() if key in sig.parameters} | |
| return fn(*args, **kwargs) |
1 participant
PR type
PR information
This PR fixes padding-free and sequence-parallel compatibility issues with newer Transformers versions.
Main changes:
cu_seq_lens_q.cache_positionto newer forward methods.This fixes issues where SDPA training under SP could see local shard length in mask creation, causing incorrect loss compared with
flash_attention_2.Experiment results
Validation performed: