adds a Transformers module#279
Conversation
📚 Documentation preview 🚀Preview URL: https://EarthyScience.github.io/EasyHybrid.jl/previews/PR279/ Note The preview will be available once the documentation build completes successfully, and will reflect the last successful build for this PR. |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive EasyHybrid.Transformers module featuring state-of-the-art architectures like TransformerModel, VisionTransformer, EncoderDecoderModel, and VisionEncoderDecoderModel, along with modern building blocks such as Grouped Query Attention (GQA), SwiGLU FeedForward, RMSNorm, and Rotary Positional Embeddings (RoPE). The review feedback identifies several critical issues that need to be addressed, primarily concerning GPU compatibility (such as CPU-bound mask generation and CPU ranges causing device mismatches), alignment of RoPE application between training and autoregressive inference, slicing of positional frequencies to prevent dimension mismatches, avoiding concrete Float32 type assertions on keyword arguments, and optimizing memory by removing unused parameter allocations and type parameters.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive EasyHybrid.Transformers module featuring modern architectures such as TransformerModel, VisionTransformer, EncoderDecoderModel, and VisionEncoderDecoderModel, along with building blocks like Grouped Query Attention (GQA), SwiGLU FeedForward, RMSNorm, and Rotary Positional Embeddings (RoPE). It also adds sequence-to-sequence data preprocessing utilities (split_seq2seq), documentation, and extensive tests. The review feedback highlights critical issues regarding GPU compatibility and performance. Specifically, when use_rope is enabled in VisionTransformer and VisionToVisionModel, the rotary positional embeddings are never computed or passed to the blocks. Additionally, several operations—such as allocating arrays with zeros, copying CPU ranges, and precomputing RoPE frequencies—are hardcoded to the CPU, which will break GPU pipelines or cause slow host-to-device transfers. The reviewer recommends using similar and device-agnostic broadcasting to ensure all computations and allocations occur on the correct device.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Some years ago I started https://github.com/lazarusA/Evie.jl, and recently I added Transformer blocks based on Lux in the Whisper package therein. And I have been thinking why just not port some of that code here? Hence this PR, which takes a lot of the conventions and utils defined there.
This PR sets the initial steps for the bigger goal of having a fully
Soils, Land, Atmosphere, Ocean, and whatnotFoundation model (+hybrid) in Julia.Next,
zarr-writerfor the model's weights, which will be useful when having global inputs and outputs (massive datasets).Some reading