From 9b91530911045f7e07f403b95ded11e2d0ad4e61 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Sun, 23 Aug 2026 22:51:42 +0000 Subject: [PATCH] Stop post-training tripping on config a model never set 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. --- src/maxtext/layers/nnx_decoders.py | 2 +- tests/utils/forward_pass_logit_checker.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1a9fdd48b0..bfb942d85d 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -438,7 +438,7 @@ def __init__( self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 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: self.hc_head = mhc.DeepSeek4HyperHead( config=config, mesh=self.mesh, diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index a51b23980f..6b79e0ae02 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -379,6 +379,12 @@ def main(config, test_args): # pylint: disable=W0621 try: max_logging.log(f"Loading tokenizer from {path}.") tokenizer = AutoTokenizer.from_pretrained(path, token=hf_token, trust_remote_code=test_args.trust_remote_code) + # Mistral and Llama tokenizers ship no pad token, and the calls below pass padding=True, which + # transformers refuses without one ("Asking to pad but the tokenizer does not have a padding + # token"). Prompts are tokenized one at a time here, so nothing is ever actually padded and + # borrowing eos changes no result -- it only satisfies the check that rejects the call. + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token break except Exception as e: # pylint: disable=broad-except,broad-exception-caught last_exception = e