Repository navigation
Conversation
5b608a7 to
35a858c
Compare
e9445a0 to
c0a093d
Compare
35a858c to
52950e3
Compare
4c4c8a5 to
a88062e
Compare
7768bf6 to
82862c3
Compare
a88062e to
e38a6d9
Compare
82862c3 to
d655d4d
Compare
e38a6d9 to
2158d8c
Compare
d655d4d to
da62987
Compare
2158d8c to
dbdd171
Compare
da62987 to
5d10546
Compare
5d10546 to
d429a6a
Compare
RyanPCo
left a comment
RyanPCo
left a comment
There was a problem hiding this comment.
Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 2 low, 5 nit.
| # -------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _qwen_snapshot() -> str | None: |
There was a problem hiding this comment.
low (test-quality): _qwen_snapshot, QWEN_SNAPSHOT/requires_qwen and _cross_attn_specs are duplicated verbatim in test_text_encoder_fixes.py. Move them, plus a Qwen snapshot fixture, into tests/unit/conftest.py. The fixture could resolve the path with snapshot_download(..., local_files_only=True) inside try/except instead of hand-globbing the HF cache layout. Also, os.environ.setdefault("HF_HUB_OFFLINE", "1") in the module-scoped fixtures is never undone and leaks into later tests in the same process. Use with pytest.MonkeyPatch.context() as mp: mp.setenv(...) with a yield fixture.
| def pretrained_reference_state_dict(self) -> Optional[dict]: | ||
| """The snapshot's own safetensors, cast exactly the way | ||
| ``from_pretrained(torch_dtype=...)`` casts them.""" | ||
| if not os.path.isdir(self._snapshot_dir): |
There was a problem hiding this comment.
low (error-handling): When HF is online, _snapshot_dir is the repo id ('Qwen/Qwen3-Embedding-0.6B'), so this returns None and verify_pretrained_weights skips the encoder silently. The 'on by default' check therefore only runs offline. Resolve the cached snapshot regardless of offline mode, for example snapshot_download(model_name, local_files_only=True) (the files are cached after from_pretrained), and log a warning in verify_pretrained_weights whenever a PretrainedWeights module is skipped.
| entry = self._text_feature_cache.get(prompt) | ||
| # Entries computed on another device are stale (the module moved): | ||
| # drop them rather than silently copy across devices. | ||
| if entry is not None and entry.device != device: |
There was a problem hiding this comment.
nit (simplification): The _apply override already clears the cache on every device or dtype move, so this per-entry entry.device != device check and its comment cannot trigger. Remove it and rely on _apply.
| compared with that reference once the model is fully built. | ||
| """ | ||
|
|
||
| _hpt_pretrained = True |
There was a problem hiding this comment.
nit (dead-code): _hpt_pretrained = True is never read; detection goes through pretrained_submodules and pretrained_reference_state_dict. Remove it. Since the mixin now exists, pretrained_module_ids and verify_pretrained_weights can use isinstance(module, PretrainedWeights) instead of duck-typed getattr.
| ) | ||
| live_hash = hash_state_dict(current, dtype) | ||
| ref_hash = hash_state_dict(reference, dtype) | ||
| label = name or root.__class__.__name__ |
There was a problem hiding this comment.
nit (simplification): Compute label at the top of the loop body and use it in both raises. The "This runs during HPT.init, before any checkpoint restore" sentence couples a generic helper to one caller, so move it to a comment at the call site in hpt.py.
| # Pretrained-weight protection | ||
| # | ||
| # ``HPTModel.finalize_modules`` runs a xavier pass (``_init_weights``) over the | ||
| # whole model once the stems / heads are ModuleDicts. That pass used to visit |
There was a problem hiding this comment.
nit (comments): This banner narrates the bug that was fixed ("That pass used to visit...", "without touching finalize_modules again"); that history belongs in the commit message. A one-line description is enough, e.g. "Modules owning pretrained weights declare them so init passes skip them and the weights can be hash-verified". Likewise, the ~1.1 GB / ~3 s figures in hpt.py and the per-entry size table in the text_cache_max_entries docstring can shrink to an order-of-magnitude note.
| its size follows the prompt: 4-25 KB for a 2-12 token prompt at | ||
| 1024-dim fp16, up to max_length * D * 2 bytes (262 KB at | ||
| max_length 128). The 8192 default is therefore ~100-200 MB of | ||
| device memory on real annotations and 2 GB in the worst case. |
There was a problem hiding this comment.
nit (documentation): nets.float() upcasts the encoder, so cached rows are fp32 and the per-entry and total figures here (and in the YAML comments) are 2x too low: about 4 GB in the worst case at 8192 entries x 128 tokens x 1024 dims. Either fix the numbers, or store entries as .to(self._load_dtype) at the cache write and cast back on read to keep the documented size.
… padding `HPTModel.finalize_modules` ran `self.apply(self._init_weights)` over the whole model, which xavier-randomised the pretrained HF encoders inside the text stems (the frozen Qwen3-Embedding weights were overwritten in place and their biases zeroed). Modules that own pretrained weights now declare them through the `PretrainedWeights` mixin, and the init pass walks the tree without ever entering such a subtree, so the ResNet backbone, `T5Encoder` and any future pretrained stem are covered by the same mechanism. `HPT.__init__` also called `finalize_modules` twice; the first call is gone. Once the model is built, moved and upcast, `verify_pretrained_weights` hashes every declared pretrained submodule and compares it with the checkpoint read straight off the local HF snapshot, failing the run on a mismatch. It costs a second full copy of the encoder (~1.1 GB fp16 at 0.6B) and ~3 s of hashing per rank, so `verify_pretrained: false` turns it off; it is on by default, because a silently re-initialised encoder trains to a plausible wrong result. Only stems that own an HF snapshot are checked -- `ResNet` returns no reference, so it is protected from the init pass but not hash-verified, which the mixin docstring now says. Finally, `QwenPerTokenEncoder.compute_latent` now passes the tokenizer attention mask into its cross-attention. Padded positions are zeroed before the biased `proj`, so they were the learned constant `proj.bias`, not zero, and got real attention weight: text-conditioned results from before this commit are not comparable with results after it. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
The annotation vocabulary is tiny next to the number of training steps, so a frozen text encoder re-runs the same 600M-parameter forward every step. `_Qwen3BaseEncoder` now keeps a lazy in-memory cache of the per-token hidden states keyed by the exact prompt string: misses go through the tokenizer + encoder in ONE sub-batch of unique strings, their un-padded (L_i, D) rows are stored detached, and the batch is re-assembled with the same LEFT padding the tokenizer produces, so `forward_with_mask` / the pooled path are unchanged. Only legal when the encoder is frozen - a trainable encoder must stay online. Knobs on the stem: `cache_text_features` (default true) and `text_cache_max_entries` (default 8192, LRU); both added to the model yamls with a Qwen stem. The cap counts entries, and an entry is the prompt's un-padded (L, D) states -- 4-25 KB for real annotations, 262 KB at max_length 128 -- so the default is ~100-200 MB of device memory in practice and 2 GB in the worst case; the docstring says so next to the knob. A batch carrying more unique prompts than the cap reads its rows from the miss pass rather than from the cache, because the LRU trim runs first and would otherwise evict entries the very next line looks up. `_apply` (to/cuda/float) clears the cache, and `clear_text_feature_cache()` / `text_feature_cache_stats()` are exposed. Padded positions come back as zeros instead of the encoder's pad-token output; every consumer masks them out. Cached-vs-online on 50 real prompts (fp16, each warmed alone: worst-case padding mismatch): max |diff| 3.6e-3 on the latents and 9.4e-2 on raw hidden states, exactly equal to the online path's own batch-padding noise (9.4e-2); in fp32 the same comparison is 1.1e-6 / 1.6e-4. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
d429a6a to
8768ba7
Compare
|
Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed. Fixed
Partial
|
HPTModel.finalize_modulesranself.apply(self._init_weights)over the wholemodel, which xavier-randomised the pretrained HF encoders inside the text stems
(the frozen Qwen3-Embedding weights were overwritten in place and their biases
zeroed). Modules that own pretrained weights now declare them through the
PretrainedWeightsmixin, and the init pass walks the tree without everentering such a subtree, so the ResNet backbone,
T5Encoderand any futurepretrained stem are covered by the same mechanism.
HPT.__init__also calledfinalize_modulestwice; the first call is gone.Once the model is built, moved and upcast,
verify_pretrained_weightshashesevery declared pretrained submodule and compares it with the checkpoint read
straight off the local HF snapshot, failing the run on a mismatch.
Finally,
QwenPerTokenEncoder.compute_latentnow passes the tokenizerattention mask into its cross-attention. Padded positions are zeroed before the
biased
proj, so they were the learned constantproj.bias, not zero, and gotreal attention weight: text-conditioned results from before this commit are not
comparable with results after it.
Folded in (to keep the stack under Graphite's 50-branch cap):
🤖 Generated with Claude Code
https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw