feat(retrieval): context aware reranker training - #3786
Open
sahel-sh wants to merge 25 commits into
Open
Conversation
| # pushing them onto a coarse part of the bf16 grid -- at magnitude ~32 the | ||
| # spacing is 0.25, which distorts the softmax and every gradient through it. | ||
| # Staying in fp32 through cross_entropy costs nothing at this tensor size. | ||
| outputs.logits = outputs.logits.view(-1, self.train_n_passages).float() / self.temperature |
Contributor
There was a problem hiding this comment.
This is a good addition, upcasting lower precision scores to fp32 is a good practice before softmax cross entropy.
This is helpful for training cross encoders in general.
| last_idx = _last_token_indices(attention_mask) | ||
| last_hidden = hidden_states[torch.arange(batch_size, device=hidden_states.device), last_idx] | ||
|
|
||
| last_logits = self.lm_head(last_hidden) # [B, vocab] |
Contributor
There was a problem hiding this comment.
Here the LM head scores all 151669 tokens. Wouldn't it be more efficient to ajust the MLP head to keep only the "yes" and "no" token outputs, before forwarding last_hidden into it, in order to save compute during training and inference?
That efficient implementation would be something like the "Feature Extraction + Pooling + Dense" from sentence-transformers (see here and here)
Contributor
Author
There was a problem hiding this comment.
Yes, full head with no slicing is wasteful for sure, both at training and at inference. My main motivation here was to keep everything identical to the base Reranking model, Qwen3-reranker-4b: no custom forward(), no custom HF registry entry, no custom inference code.
Contributor
|
/ok to test 910cfae |
svcnvidia-nemo-ci
added
the
waiting-on-customer
sahel-sh
added a commit
to sahel-sh/Automodel
that referenced
this pull request
Sep 3, 2026
…nd model coverage Three review findings on NVIDIA-NeMo#3786. Declare tie_word_embeddings_support = TieSupport.BOTH on Qwen3RerankerForCausalReranking and implement a model-local tie_weights(). The published checkpoints set tie_word_embeddings=True, and the class reuses the backbone's lm_head rather than adding its own, so both layouts load. The tie has to be applied explicitly because transformers v5 does not reliably tie a custom model from the dict-shaped _tied_weights_keys alone -- and it matters here because scoring reads the yes/no rows of lm_head, so a tie that silently failed would leave the head drifting from the embeddings it shares and the published checkpoint's scores would not reproduce. Adds tests for the tied alias, the untied case, and the declared policy. Switch test_build_generic_hf_model_score_task's fake model_type from "qwen3" to "bert". Registering qwen3 as a score backbone meant that fixture had stopped exercising the generic AutoModelForSequenceClassification fallback and was instead selecting the reranker, then failing to load weights from an empty directory. Add Qwen3-Reranker-4B to the reranker coverage page with the literal architecture name, and update the release-discovery expectation to include it. The introduction no longer claims only Llama has an optimized path, and the "Optimized Backbones (Bidirectional Attention)" heading is now just "Optimized Backbones", since this backbone stays causal. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
sahel-sh
force-pushed
the
retrieval_context_aware_reranker
branch
from
September 3, 2026 13:37
f2c32f0 to
43f814e
Compare
Contributor
|
🌿 Preview your docs: https://nvidia-preview-preview-81ea469e2416.docs.buildwithfern.com/nemo/automodel |
jgerh
reviewed
Sep 3, 2026
jgerh
left a comment
jgerh
left a comment
Contributor
There was a problem hiding this comment.
Completed tech pubs review of docs/model-coverage/reranker/index.mdx and provided a few copyedits.
Adds the pieces needed to train a Qwen3 reranker on data where each query
carries context -- the reasoning trace that produced it and the originating
question -- rather than the query alone.
Qwen3RerankerForCausalReranking keeps the backbone causal and adds no
parameters: it scores a pair as logit("yes") - logit("no") at the final
position, which is how Qwen/Qwen3-Reranker-* is trained and served. Without it
a Qwen3 checkpoint loaded for scoring falls through to a sequence-classification
head initialised from noise, discarding the relevance signal the checkpoint
already holds. Checkpoints serialize as plain Qwen3ForCausalLM so they load in
vLLM without trust_remote_code.
ContextAwareRerankerCollator embeds the context fields in the <Query> block and
selects the instruction from whichever fields survive its per-query dropout, so
one model covers all four prompt modes and the instruction always describes the
prompt actually sent. With no context fields the prompt is byte-identical to the
stock Qwen3-Reranker one. Chat markers are parameters, defaulting to ChatML.
make_context_aware_retrieval_dataset carries those fields through the transform
and holds out validation at the level of a group key, so rows sharing a question
cannot land on opposite sides of the split -- a dataset holding two labelings of
the same query leaks badly otherwise.
The cross-encoder loss now scales logits by the recipe-level temperature, as the
bi-encoder path does. Temperature belongs to the training objective, not to the
model, so it stays out of the checkpoint.
Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
The cross-encoder backbone emits bf16 logits, and the recipe divided them by the temperature while still in bf16. At temperature 0.1 that magnifies them 10x onto a coarse part of the bf16 grid -- at magnitude ~32 the representable spacing is 0.25 -- which distorts the softmax and every gradient flowing through it. This was not theoretical. v60 failed to reproduce v50 on the ported code path: it trained the same distance from the base model but in a different direction and scored worse on BRIGHT and BCP. Every v60/v61 loss lands exactly on the bf16 grid (0.3086, 0.3926, 0.4883) where v50/v51's do not. Up-casting before the division fixes it, and costs nothing at this tensor size. The same change is applied to the validation path so train and val losses stay comparable. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…obal-query context The example config was the pre-global-query recipe: it declared only reasoning_column and reasoning_drop_prob, split validation on query_id, and pointed at a JSONL with no global_query field, so the four-mode context dropout the collator supports could not be exercised from it. Bring it in line with the configs actually in use: temperature 0.1 -> 0.3 num_epochs 2 -> 1 global_query_column added global_query_drop_prob added, 0.5 on train and validation reasoning_drop_prob validation 0.0 -> 0.5 validation_group_key query_id -> global_query data_dir_list agentir_inline -> agentir_v39_inline The data file has to move with the settings: agentir_inline.jsonl carries no global_query, so the new probabilities would have been silently inert against it. Validation mirrors the training probabilities rather than disabling dropout. Validation should measure the same prompt-mode distribution the model trains on; drop_seed is fixed, so the mix is identical on every pass regardless. Grouping the split on global_query keeps every sub-query of one originating question on the same side. Keying on query_id put siblings on opposite sides, which leaks. Replace the note on how the JSONL was produced with the schema the loader expects -- required and optional keys, how pos_doc/neg_doc map onto n_passages, and that a missing or blank reasoning/global_query is treated as absent rather than empty. Adds a TODO to source the data from the published Hugging Face dataset instead of a local path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
retrieval_collator: validate instructions at construction. _normalize_instructions accepted any field names and _format_one fell back to the no-context instruction when a surviving combination had no entry, so a misspelled or omitted key trained a context-bearing prompt under the base instruction -- invisible in the loss, visible only as a model that ignores its context at eval. Unknown and missing modes are now both errors, raised before the first batch rather than mid-run, and the fallback is gone. retrieval_dataset_inline: reject missing validation_group_key values before sorting. load_datasets fills absent extra columns with None, so a key present on some rows and absent on others reached sorted() as a mix of str and None and raised TypeError from inside dataset construction, reading as a library bug rather than a data problem. Blank strings are rejected on the same path since they group just as badly. qwen3_reranker/__init__: guard the public re-export with safe_import_from. A hard import made the package unimportable when the Qwen3 stack is unavailable and took the parent package with it; the placeholder now raises only at point of use, naming the missing dependency. qwen3_reranker/model: pass return_dict=True to the decoder. last_hidden_state is only reachable on a ModelOutput and the call did not forward the wrapper's return_dict, so config.return_dict=False made the decoder hand back a tuple and the attribute access raise. Confirmed directly: without the pin the decoder returns a plain tuple with no last_hidden_state. qwen3_reranker/model: preserve rope_scaling when rewriting rope_parameters. The scaling config lives nested inside rope_parameters; popping that dict and then defaulting rope_scaling to None discarded a non-default RoPE scale, so a saved checkpoint reloaded in HF or vLLM with different position encoding than it trained with. Verified a yarn factor survives the export while an unscaled config still emits None. Also bumps the copyright year on the one file whose header was touched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…ry text The group-aware split kept every sub-query of one research session on the same side by comparing the global_query TEXT. That is fragile: the key is a long free-text field, so any drift between two rows of the same session splits it across the train/validation boundary -- exactly the leak the grouping exists to prevent. It is not hypothetical. In the ORBIT trajectories one session carries two different global_query strings, because the agent rephrased the question mid-run: "What is the number of points in the complete graph used for a pencil-and-paper game ..." "What is the outcome with perfect play in a combinatorial game where two players ..." Grouping on text therefore made one session into two groups and, symmetrically, merged two other sessions that happened to share a string. The datasets now carry an explicit global_qid column -- the session id, shared by every turn of a trajectory and by every labeling variant of the same question -- and the example config groups on it. Matching an id is exact and cannot drift. This also matters for the mixes, where query_id is NOT unique: mix_nemotron_v39_v48 has each of its 5,238 sub-queries under three different labelings, and global_qid keeps all three on one side. Versioned run configs are deliberately left on global_query. They record experiments that already trained under that key, and the two keys select materially different partitions (only 150 rows overlap on mix_nemotron_v39_v48), so rewriting them would both misdescribe finished runs and silently change the split for any rerun. Also fail clearly when the split key is not a column of the dataset. It previously surfaced as a KeyError from inside dataset construction, which reads as a library bug rather than a configuration mistake; the error now names the key and lists the available columns. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…right years The example config groups the validation split on global_qid but never listed the field in the schema block above it, so a reader building a dataset from that contract would omit the one column the split depends on. Document it as required, and say what it means: the session id, shared by every turn of a trajectory and by every labeling variant of the same question. Also expand the note on why the split keys on it. query_id is not merely a worse choice, it is unusable in the mixes -- mix_nemotron_v39_v48 carries each sub-query three times, once per labeling -- and global_query text is a long free-text field that splits a session in two when its rows disagree by a character. Copyright headers on the files touched in this branch move to 2026. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
The header year records when a file was created, not when it was last edited. The previous commit bumped retrieval_collator.py and retrieval_dataset_inline.py to 2026 simply because this branch modifies them; both already exist on internal and keep 2025. Only the files this branch actually adds carry 2026: qwen3_reranker/model.py, qwen3_reranker/__init__.py and the qwen3_4b_reranker_agentir.yaml example. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
forward() declares return_dict: Optional[bool] = None, and the final branch tested that value directly. A caller that omitted the argument therefore hit `if not return_dict` with None, took the tuple path, and never consulted config.use_return_dict -- so a model configured to return ModelOutput silently handed back a bare tuple and callers reading SequenceClassifierOutputWithPast.logits broke. Resolve it once, at the top of forward, the same way qwen2/model.py, qwen3/model.py and llama_nemotron_vl/model.py do, then branch on the resolved value. This is separate from the return_dict=True pinned on the inner decoder call: that one exists because last_hidden_state is only reachable on a ModelOutput, and is independent of what this wrapper returns. The comments now say so explicitly, since the two were easy to conflate. Verified across both config settings: with config.return_dict True an omitted argument now yields SequenceClassifierOutputWithPast (previously a tuple); with it False an omitted argument yields a tuple; and an explicit True/False is honoured either way. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Train and validation are two SEPARATE _group_aware_split calls, from two separate configs. They were complementary only because both happened to see the same groups in the same order with the same seed: the implementation materialised the group list and shuffled it, so any difference in file order -- or one extra row on one side -- permuted the shuffle and moved groups across the boundary. That reintroduces exactly the leakage the grouping exists to prevent, silently, with no error raised. Assign each group by hashing (seed, group) instead. A group's side is now a pure function of its own id, so the two calls agree whatever order the rows arrive in. Select by RANK on the hash rather than thresholding the score. Thresholding is simpler but makes the validation size binomial: 0.2 of 40 groups came out as 7 rather than 8 in testing, and a 10% split of 249 groups would vary by roughly +/-5, where the previous implementation always produced exactly round(n * fraction). Ranking keeps that exact count. Verified: train/val disjoint and covering; identical split under a shuffled file order (the property the old code lacked); exact counts (25 of 249, 28 of 276, 80 of 800 at fraction 0.1); a different seed still yields a different split; and adding 7 new groups left all 8 original validation groups in validation, where the old shuffle would have repermuted everything. Also log the seed, fraction and a fingerprint of the chosen validation set. The train and validation builds emit one line each, back to back, so a seed or fraction mismatch between the two configs shows up as two differing fingerprints instead of as silent leakage. Separately, complete the truncated sentence from e1a12417. Rows with fewer than seven negatives are not dropped: _retrieval_transform_func indexes negatives modulo the list length, so they are CYCLED to fill the group, and only a row with no negatives at all raises. The original "dropped by the loader" claim was wrong. Note this changes the split for existing datasets, so runs trained before this commit are not directly comparable to later ones. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Follow the existing StepScheduler.set_epoch pattern instead of injecting a closure from the recipe. set_epoch already forwards the epoch to the sampler; forward it to the dataloader's collate function too, and give ContextAwareRerankerCollator a set_epoch method. _epoch_fn and set_epoch_source are gone, along with the lambda in train_bi_encoder. The epoch is held in a shared-memory tensor rather than a plain int because collate_fn runs inside the DataLoader workers. Under persistent_workers=True those workers outlive the epoch boundary, so an attribute assigned in the parent would never reach the already-forked children and every epoch would replay epoch 0's drop mask. Train-only behaviour is now structural: the validation loader is built separately and is not owned by the scheduler, so it stays at epoch 0 and keeps a fixed prompt mix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Hashing (seed, group) made a group's SCORE independent of load order, but
the selection built on top of it did not inherit that property: n_val came
from len(groups) and membership came from a group's rank among all others,
so the cutoff was a property of the whole group set.
Two consequences, both silent:
* Adding groups can evict an existing group from validation into
training, so a held-out group becomes training data when the corpus
grows and validation sets are not nested across dataset versions.
* Train and validation are built by two separate calls from two separate
configs. If those do not resolve to an identical group set, the calls
compute different cutoffs and a group can land in training on one side
and validation on the other.
Compare the score against the fraction instead of ranking, so a group's
side depends on nothing but (seed, group, fraction). This costs an exact
count -- the size is binomial rather than round(n * fraction) -- which is
a cheaper failure than leakage.
Also sort the hashlib import that the previous commit left out of order.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
A thresholded split has a binomial size, so a dataset with few groups can draw every group onto one side even though validation_fraction > 0. Both outcomes are unusable and both surface far from the cause: an empty validation side yields a zero-length eval dataset, and an empty training side yields a run with nothing to train on. Raise where the seed, fraction and group key that produced it are still in scope, and name them in the message. validation_fraction=0 is untouched: asking for no held-out slice stays valid. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Qwen3-Reranker is a causal LM: its published config.json declares architectures ["Qwen3ForCausalLM"], the model card loads it with AutoModelForCausalLM, and the official compute_logits scores by reading full-vocabulary next-token logits at the final position and indexing the yes/no ids. Reranking is next-token prediction here, not classification, so Qwen3ForCausalLM is the faithful base class: it keeps that identity, reuses upstream's model + lm_head construction and weight tying rather than duplicating it, and keeps state-dict keys identical to the backbone. Saved checkpoints deserialize as plain Qwen3ForCausalLM, so a finetuned model scores through the stock HF and vLLM paths exactly like the backbone. Document that, and address the real concern behind the review: the training-time forward returns [batch, 1] scores, so a generation caller asking for per-position logits was silently handed the wrong shape. Reject a non-zero logits_to_keep by name instead; 0 and None remain accepted. Tests assert sigmoid(score) equals the model card's compute_logits p(yes) under both padding sides, plus the weight-key and config round-trip contracts that keep checkpoints loadable as stock Qwen3ForCausalLM. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…oken index Two review follow-ups in the same forward path. Use transformers' @can_return_tuple instead of hand-rolling return_dict: forward always builds SequenceClassifierOutputWithPast and the decorator converts it via to_tuple() when return_dict=False is passed or set on the config. Slightly more standard than before -- to_tuple() emits every non-None field in declaration order, so loss still comes first when present, where the manual branch emitted only (loss, logits). Make _last_token_indices branchless. Choosing the padding side from a tensor value forced a device-to-host sync on CUDA and a graph break under torch.compile. Masking position indices and taking the row max needs neither, and it also fixes a latent bug: the old whole-batch padding-side guess mislabelled rows in a batch that mixes left- and right-padded sequences (returning [2, 2] where [4, 2] is correct), while the masked-max resolves each row independently. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…del package Every constant this collator carries is part of the Qwen3-Reranker prompt contract -- the ChatML markers, the empty <think> block that makes the final tokens a yes/no next-token prediction, the model card's system message, and the yes/no label semantics. That is model-specific logic, which belongs under components/models/<model>/ rather than the generic dataset collators. Rename to Qwen3ContextAwareRerankerCollator and base it directly on DataCollatorWithPadding. It overrode both methods CrossEncoderCollator defines and used neither prompt_template nor format_text, so the old base was vestigial; standing alone costs no duplication and drops the dependency on the generic retrieval collator module. The class is new on this branch and absent from origin/internal, so no released _target_ path changes. Config references are updated in place. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…eckpoint round trip Both ends of validation_fraction failed silently. The split runs under `if validation_fraction > 0`, so a negative value skipped it entirely and the validation build returned the whole dataset -- every training row also an eval row, with no error raised. A value of 1 or more did raise, but from the empty-side check, whose message blames the group count and suggests changing the seed, none of which is the actual problem. Reject the value in the builder, before any data is read. Separately, the config identity rewrite was only asserted through to_dict(), which leaves save_pretrained/from_pretrained untested: a weight-key rename, a dropped config field, or a tie_word_embeddings mismatch would all pass that check and still yield a checkpoint that reloads wrong. Add three round-trip tests -- reload as stock Qwen3ForCausalLM and score it the model-card way, reload as the reranker and confirm the yes/no ids survive, and confirm the stock reload's state-dict keys are exactly the reranker's. Both new tests were mutation-checked: reverting the architectures rewrite fails the stock-reload test and dropping the yes/no ids on save fails the reranker-reload test, in each case beyond the pre-existing to_dict assertion. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…elForCausalLM The claim these tests exist to defend is that a saved checkpoint loads through the stock path with no custom code and no trust_remote_code. Reaching for Qwen3ForCausalLM directly bypasses that path: the concrete class never consults model_type or architectures, so it loads the weights happily even when the serialized identity still names this package's class -- the one case that would actually send a consumer looking for custom code. All the real coverage came from a separate string assertion on config.json rather than from the load. Both round-trip tests now go through AutoModelForCausalLM, and the stock-reload test asserts the resolution rather than only the loadability: the checkpoint must land on upstream's Qwen3ForCausalLM and its config must come back as Qwen3Config, not the reranker subclass. This turns the identity rewrite into something the load can fail on. Mutation-checked: serializing model_type as "qwen3_reranker" now fails both round-trip tests, where previously nothing but the config.json string assertion noticed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Upstream lint has tightened since this code was written on the internal branch: UP045 now rejects Optional[X] in favour of X | None. Mechanical, no behaviour change. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
Registering qwen3 with a "score"-only task map removed qwen3 embedding
backbones from the generic HuggingFace path. Before the entry existed,
SUPPORTED_BACKBONES.get("qwen3") returned None and the caller fell through
to AutoModel with is_causal=False; afterwards the same request raised
ValueError: Unsupported task 'embedding' for model type 'qwen3'.
Restore the fallback on both paths. _get_supported_backbone_class returns
None for "embedding" when the task map has no entry, and the guard in
_build_backbone_from_extracted_submodel now allows "embedding" alongside
the existing "score" exception, since both have a generic fallback below.
Genuinely unknown tasks still raise, so misconfiguration is not masked.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
…nd model coverage Three review findings on NVIDIA-NeMo#3786. Declare tie_word_embeddings_support = TieSupport.BOTH on Qwen3RerankerForCausalReranking and implement a model-local tie_weights(). The published checkpoints set tie_word_embeddings=True, and the class reuses the backbone's lm_head rather than adding its own, so both layouts load. The tie has to be applied explicitly because transformers v5 does not reliably tie a custom model from the dict-shaped _tied_weights_keys alone -- and it matters here because scoring reads the yes/no rows of lm_head, so a tie that silently failed would leave the head drifting from the embeddings it shares and the published checkpoint's scores would not reproduce. Adds tests for the tied alias, the untied case, and the declared policy. Switch test_build_generic_hf_model_score_task's fake model_type from "qwen3" to "bert". Registering qwen3 as a score backbone meant that fixture had stopped exercising the generic AutoModelForSequenceClassification fallback and was instead selecting the reranker, then failing to load weights from an empty directory. Add Qwen3-Reranker-4B to the reranker coverage page with the literal architecture name, and update the release-discovery expectation to include it. The introduction no longer claims only Llama has an optimized path, and the "Optimized Backbones (Bidirectional Attention)" heading is now just "Optimized Backbones", since this backbone stays causal. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Follow-up on the review: declare _tied_weights_keys explicitly and call reject_unsupported_tie_word_embeddings() at the top of __init__, on the original config and before super().__init__, so an unsupported tie_word_embeddings fails at construction rather than after weights have loaded. BOTH accepts either setting, so the rejection never fires for this class today. It is here so that narrowing the policy later cannot silently produce a model with a randomly-initialised head -- which for this class would be silent rather than obvious, since scoring reads only the yes/no rows of lm_head. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The example yaml linter requires every example wandb: block to set enable: false, so a user who copies a recipe does not start logging to someone else's project by accident; enabling is a deliberate flip to true. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Applies edit comments Co-authored-by: jgerh <163925524+jgerh@users.noreply.github.com> Signed-off-by: Sahel Sharifymoghaddam <sahel.sharifi@gmail.com>
… columns codecov/patch reported 67.63% of the diff hit against a target of 80%, the only failing check on the PR. The uncovered code was not incidental: the collator's prompt assembly and the dataset's context-column wiring are the substance of this change, and neither had a test. Collator, now 100% of the added lines. Covers the four prompt modes with the drop probabilities pinned to 0.0 and 1.0 so a mode is selected rather than sampled: that base mode carries no markers and stays byte-identical to the stock Qwen3-Reranker prompt, that the instruction always matches the fields which survived the draw, that per-field caps apply before assembly so a long trace cannot starve the query, and that the chat markers survive truncation. Also the instructions= normalisation, including tuple and comma-separated string keys for YAML, and the errors for an unknown or missing mode. Model, now 100%. Adds the serialization of a non-default RoPE scale, which would otherwise reload with different position encoding than it trained under; scoring with unset yes/no ids raising rather than returning a constant; from_pretrained recovering the ids from the tokenizer for a checkpoint saved as a plain causal LM; and the pointwise labels path the recipe does not use. Dataset, 54.8% to 92.9%. Covers context columns repeating once per document so every row of a listwise group carries the same context, a configured column absent from the data being skipped, the builder threading the column names through to the loader and renaming them to what the collator reads, and the argument validation. Patch coverage measured locally over the added lines goes from 64.90% to 97.70%. All new tests are CPU-only unit tests per the testing skill, and assert on observable behaviour and failure modes rather than mirroring the implementation. Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com>
sahel-sh
force-pushed
the
retrieval_context_aware_reranker
branch
from
September 4, 2026 21:55
81ea469 to
0fc5662
Compare
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
5 participants
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.
What does this PR do ?
Adds context-aware cross-encoder reranker training: finetunes a causal LM as a listwise
reranker that scores a query–passage pair using the context the query came from — the
reasoning trace that motivated it and the originating question — not the query alone.
One row per sub-query, 1 positive and 7 hard negatives, one softmax per group. The collator
drops each context field independently, so a single model covers all four prompt shapes.
The model adds no parameters: it scores
logit("yes") - logit("no")off the pretrainedlm_head, and checkpoints serialize as a plainQwen3ForCausalLMthat loads in stock HF andvLLM without
trust_remote_code.Changelog
models/qwen3_reranker/model.py:Qwen3RerankerForCausalReranking; raw yes/no log-odds,padding-side agnostic, serialized identity rewritten to stock
Qwen3ForCausalLM.models/qwen3_reranker/collator.py: per-field context dropout; the<Instruct>line ischosen from whichever fields survive, and per-field token budgets are applied before
assembly so a long trace cannot crowd out the query.
datasets/llm/retrieval_dataset_inline.py: context columns and a group-aware validationsplit keyed on a hash of
(seed, group), so independently-built train/val sets staycomplementary regardless of load order.
recipes/retrieval/train_cross_encoder.py: listwise cross-entropy, cast to fp32 before thetemperature division — these recipes run native bf16, so the loss would otherwise quantize.
_transformers/{registry,retrieval}.py: register the reranker forscore, keeping thegeneric
AutoModelfallback forembedding.training/step_scheduler.py,recipes/retrieval/train_bi_encoder.py: forward the epoch tothe training collator only, so validation keeps a fixed prompt mix.
examples/retrieval/cross_encoder/qwen3_4b_reranker_agentir.yaml: example recipe with thedata schema documented inline.
Before your PR is "Ready for review"
Pre checks:
44 new tests in 4 files, plus 3 in the existing step-scheduler suite: scoring parity with the
published
compute_logitsunder both padding sides, checkpoint round-trips reloaded throughAutoModelForCausalLM, split independence from load order, and theembeddingfallback.Additional Information