Skip to main content

​
Overview

The trainer.py module extends the Hugging Face Trainer with custom implementations optimized for Qwen-VL models. It includes custom attention mechanisms using Flash Attention 2, monkey-patched forward methods for various Qwen model versions, and a specialized optimizer creation method that supports different learning rates for vision tower and multimodal projector components.

​
Custom Optimizer Creation

​
create_optimizer()

Creates an optimizer with separate learning rate configurations for different model components.
​
float
default:"None"
Learning rate for the multimodal projector (merger) module. If set, the projector parameters will use this learning rate instead of the base learning rate.
​
float
default:"None"
Learning rate for the vision tower (visual encoder) module. If set, the vision tower parameters will use this learning rate instead of the base learning rate.
​
float
Weight decay coefficient applied to parameters with decay (excludes bias terms).
Returns: Configured optimizer instance with parameter groups for:
  • Base model parameters (with and without weight decay)
  • Vision tower parameters (with and without weight decay, if vision_tower_lr is set)
  • Multimodal projector parameters (with and without weight decay, if mm_projector_lr is set)

​
Flash Attention Methods

​
flash_attention_forward()

Custom Flash Attention 2 forward pass for efficient attention computation.
​
torch.nn.Module
required
The attention module.
​
torch.Tensor
required
Query tensor with shape (batch, head, seq_len, dim).
​
torch.Tensor
required
Key tensor with shape (batch, head, seq_len, dim).
​
torch.Tensor
required
Value tensor with shape (batch, head, seq_len, dim).
​
torch.Tensor
default:"None"
Cumulative sequence lengths tensor for variable-length attention.
​
float
default:"0.0"
Dropout probability for attention weights.
​
float
default:"None"
Attention scaling factor.
​
int
default:"None"
Sliding window size for local attention.
​
float
default:"None"
Softcap value for attention logits.
Returns: Tuple of (attn_output, None) where attn_output is the attention output tensor.

​
qwen2vl_forward()

Custom forward method for Qwen2-VL and Qwen2.5-VL attention layers.
​
torch.Tensor
required
Input hidden states with shape (batch_size, seq_len, hidden_dim).
​
torch.Tensor
default:"None"
Attention mask tensor.
​
torch.LongTensor
default:"None"
Position indices for positional embeddings.
​
Cache
default:"None"
Cached key/value states for efficient generation.
​
bool
default:"False"
Whether to return attention weights.
​
bool
default:"False"
Whether to use key/value caching.
​
torch.LongTensor
default:"None"
Position indices for cache updates.
​
tuple[torch.Tensor, torch.Tensor]
default:"None"
Precomputed rotary position embeddings (cos, sin).
Returns: Tuple of (attn_output, attn_weights) where weights may be None.

​
qwen3vl_forward()

Custom forward method for Qwen3-VL and Qwen3-VL-MoE attention layers.
​
torch.Tensor
required
Input hidden states.
​
tuple[torch.Tensor, torch.Tensor]
required
Rotary position embeddings (cos, sin).
​
torch.Tensor
default:"None"
Attention mask tensor.
​
Cache
default:"None"
Cached key/value states.
​
torch.LongTensor
default:"None"
Position indices for cache updates.
Returns: Tuple of (attn_output, attn_weights).

​
Utility Functions

​
replace_qwen2_vl_attention_class()

Replaces the default attention forward methods in transformers with optimized Flash Attention implementations for all supported Qwen-VL model variants:
  • Qwen2-VL
  • Qwen2.5-VL
  • Qwen3-VL
  • Qwen3-VL-MoE
Also replaces causal mask creation functions to optimize for packed training. Prints the trainable status of vision module components. Output:
  • Trainable/non-trainable attention block indices
  • Merger module trainable status
Prints the trainable status of language model components. Output:
  • Embed tokens trainable status
  • Trainable/non-trainable decoder layer indices

​
Usage Example

​
Notes

  • Flash Attention 2 is used by default for all attention computations
  • The custom optimizer supports up to 6 parameter groups with different learning rates
  • Vision tower parameters are identified by “visual” in their name
  • Multimodal projector parameters are identified by “merger” in their name
  • Weight decay is not applied to bias parameters