Stop post-training tripping on config a model never set - #4973
Stop post-training tripping on config a model never set#4973ecnal-cienet wants to merge 1 commit into
Conversation
Two places assume a field that only some configs carry, and both fail on a model that has nothing to do with the field. The decoder reads mhc_expansion_rate directly. Every config the pre-training path builds carries it; post-training builds its models through one that does not, so a run that never touches deepseek4 still raised AttributeError on the way past the check. Reading it with the default the attribute would have held leaves the deepseek4 behaviour alone. The logit checker passes padding=True to tokenizers that ship no pad token, which Mistral's and Llama's do not, and transformers refuses outright: "Asking to pad but the tokenizer does not have a padding token". Prompts are tokenized one at a time there, so nothing is ever actually padded and borrowing eos changes no result -- it only satisfies the check that rejects the call.
![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 safe attribute retrieval for mhc_expansion_rate in nnx_decoders.py and adds a fallback pad token assignment using the EOS token in forward_pass_logit_checker.py for tokenizers that lack one. The review feedback suggests improving robustness by handling cases where mhc_expansion_rate is explicitly set to None to avoid a TypeError, and adding a fallback to bos_token in the tokenizer configuration in case eos_token is also None.
| self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL | ||
|
|
||
| if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: | ||
| if getattr(config, "mhc_expansion_rate", 1) > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: |
There was a problem hiding this comment.
If mhc_expansion_rate is explicitly set to None in the configuration, getattr(config, "mhc_expansion_rate", 1) will return None because the attribute exists but its value is None. Comparing None > 1 will then raise a TypeError. To prevent this, we should ensure a fallback value of 1 is used if the attribute is either missing or None.
| if getattr(config, "mhc_expansion_rate", 1) > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: | |
| if (getattr(config, "mhc_expansion_rate", 1) or 1) > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: |
| if tokenizer.pad_token is None: | ||
| tokenizer.pad_token = tokenizer.eos_token |
There was a problem hiding this comment.
If tokenizer.eos_token is also None, setting tokenizer.pad_token = tokenizer.eos_token will keep the pad token as None, which will still cause the padding error. Adding a fallback to tokenizer.bos_token improves robustness for tokenizers that might not have an EOS token defined but have a BOS token.
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token or tokenizer.bos_token |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
1 participant
Description
Two places read a config field that only some configs carry, and both fail on a model that has nothing to do with the field. Neither is specific to post-training, but post-training is where they surface, because it builds its models through a config that pre-training does not use.
nnx_decoders.pyreadsmhc_expansion_ratedirectly. Every config the pre-training path builds carries it; the one post-training builds does not, so a run that never touches deepseek4 still raisedAttributeErroron the way past the check. Reading it with the default the attribute would have held leaves the deepseek4 behaviour alone.forward_pass_logit_checker.pypassespadding=Trueto the tokenizer. Mistral's and Llama's ship without a pad token, and transformers refuses outright: "Asking to pad but the tokenizer does not have a padding token". Prompts are tokenized one at a time there, so nothing is ever actually padded and borrowing eos changes no result — it only satisfies the check that rejects the call.Tests
tests/unit/nnx_decoders_test.py— 47 passed,.The logit checker change has no test of its own: it is only reachable through a conversion correctness run, which needs the weights. It was exercised by hand while measuring mistral-7b against the Llama mapping (see the model registry PR), which is where the missing pad token first blocked the check.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.