[models] - feat(LayerSkip): implement LayerSkip model - #86
Merged
Conversation
…stic layer dropout for transformer layers - Introduce LayerDropoutWrapper class to selectively skip transformer layers during training based on a dropout probability - Forward pass conditionally bypasses layer execution to simulate layer dropout, retaining initial hidden states when skipped
…um callback for dynamic layer training - Introduce LayerSkipCurriculumCallback controlling layer contributions in BERT-based models - Support rotational and gradual curriculum types to adjust layer training dynamically - Allow for fixed training of the last layer through a configurable option
…out and early exit - Adds LtLayerSkip, a new Lightning module for fine-tuning with layer dropout and early exit strategy - Supports dynamic dropout scheduling and loss scaling based on model depth for efficient training - Enables both training and inference modes, with the ability to specify an early exiting layer during inference - Includes methods for forward passes with all layers or an early exit, as well as customized prediction steps
…el enhancements - Add tests for dropout schedule ensuring first layer has zero and last matches p_max - Test forward pass covers all model layers by examining tuple output and shape - Validate computed loss scales are in ascending order and sum to approximately one - Verify rotational curriculum functioning with LayerSkipCurriculumCallback in a module - Check early exit inference mode to ensure correct tensor output with specified exit layer
- Add configuration mapping for LayerSkip training setup - Include LtLayerSkip model to the available models for training and usage - Integrate LayerDropoutWrapper to the custom transformers module - Implement LayerSkipCurriculumCallback for dynamic training adjustments
JulesBelveze
marked this pull request as draft
January 31, 2026 11:18
…n Seq2SeqTransformerDataModule - Implement caching for tokenizer creation to avoid redundant loads in `_get_seq2seq_tokenizer` function - Add support for optional source and target prefixes to be applied during tokenization [bert_squeeze/models] - refactor: streamline LayerDropoutWrapper execution path - Simplify conditional layer execution in `LayerDropoutWrapper` for clearer logic - Centralize input preparation and hidden state retrieval in LtLayerSkip for evaluation steps - Extract repeated evaluation logic into a dedicated `_eval_step` method within LtLayerSkip
…for LayerSkip model - Introduce `train_layerskip.yaml` with configurations for training LayerSkip model in bert_squeeze - Include callback for curriculum learning with rotational period and optional training of the last layer - Define layer-specific learning rates and layer decay for discriminative learning approach - Set up model with parameters like exit layer, dropout schedule, and pretrained BERT model - Configure data module for processing with tokenizer and max sequence length settings
# Conflicts: # bert_squeeze/data/modules/transformer_module.py
…Skip - Enhanced the LayerDropoutWrapper to handle stochastic layer dropout more effectively, ensuring better management of hidden states and outputs during training. - Updated the LtLayerSkip class to enforce stricter validation on parameters and improved the handling of early exit logic, making the model more robust and user-friendly.
- Enhance the LayerDropoutWrapper to better manage wrapped layers and their parameters, improving encapsulation and synchronization. - Update LtLayerSkip to support input embeddings directly, allowing for more flexible input handling during inference.
- Enhanced the LayerDropoutWrapper to ensure proper type checking and state synchronization for the wrapped layer. - Simplified the forward pass logic and improved error handling for layer outputs, ensuring compatibility with transformer layers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds LayerSkip support for BERT sequence classification!
Here are the additions:
Note: the current implementation is limited to BERT models