support embedding MRL#9528
Conversation
![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 support for Matryoshka Representation Learning (MRL) in both Megatron and standard embedding trainers by adding the mrl_dims argument. This argument allows users to specify multiple embedding dimensions and their corresponding loss weights, computing and aggregating the weighted losses across these truncated dimensions. The review feedback highlights two key areas for improvement: first, validating that the configured MRL dimensions do not exceed the model's actual hidden dimension to prevent silent slicing errors; second, ensuring that the parsed mrl_dims is explicitly validated as a dictionary before accessing its items to avoid potential runtime errors.
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 Matryoshka Representation Learning (MRL) support for embedding training by adding the mrl_dims parameter to both standard and Megatron training arguments, along with corresponding documentation updates. The embedding trainers are updated to slice the last hidden state, apply L2 normalization, and compute weighted losses across the specified dimensions. The review feedback highlights critical issues that need to be addressed: an evaluation discrepancy where metrics are computed on unnormalized, full-dimensional embeddings instead of the truncated and normalized MRL dimensions; a potential crash if all configured dimensions exceed the model's hidden size; and a lack of type validation when parsing mrl_dims which could lead to runtime AttributeError crashes.
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.
3 participants
No description provided.