Add support for Qwen3-4B-Thinking-2507

[Rohan-Bierneni]

Description

Maxtext already has support for Qwen3-4B architechture, but not for the -Thinking variant out of the box. This pr aims to leverage the already implemented dense Qwen3 architechture in Maxtext to add support for the Thinking variant of Qwen3-4B.

FIXES: b/440393388

Tests

Ran forward_pass_logit_checker.py on on a converted huggingface to maxtext checkpoint via the checkpoint util and achieved a acceptable average kl-divergencevalue of 0.000514: https://paste.googleplex.com/6706050174681088

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[github-actions]

🤖 Hi @Rohan-Bierneni, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[RissyRan]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

It seems being throttle again:

Attempt 1 failed with status 429. Retrying with backoff... ApiError: {"error":{"message":"{\n  \"error\": {\n    \"code\": 429,\n    \"message\": \"You exceeded your current quota, please check your plan and billing details. For more information on this error, head to: [https://ai.google.dev/gemini-api/docs/rate-limits.\\n*](https://ai.google.dev/gemini-api/docs/rate-limits.//n*) Quota exceeded for metric: generativelanguage.googleapis.com/generate_content_free_tier_requests, limit: 2\\nPlease retry in 54.749946391s.\",\n    \"status\": \"RESOURCE_EXHAUSTED\",\n    \"details\": [\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.QuotaFailure\",\n        \"violations\": [\n          {\n            \"quotaMetric\": \"generativelanguage.googleapis.com/generate_content_free_tier_requests\",\n            \"quotaId\": \"GenerateRequestsPerMinutePerProjectPerModel-FreeTier\",\n            \"quotaDimensions\": {\n              \"location\": \"global\",\n              \"model\": \"gemini-2.5-pro\"\n            },\n            \"quotaValue\": \"2\"\n          }\n        ]\n      },\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.Help\",\n        \"links\": [\n          {\n            \"description\": \"Learn more about Gemini API quotas\",\n            \"url\": \"[https://ai.google.dev/gemini-api/docs/rate-limits\](https://ai.google.dev/gemini-api/docs/rate-limits/)"\n          }\n        ]\n      },\n      {\n        \"@type\": \"type.googleapis.com/google.rpc.RetryInfo\",\n        \"retryDelay\": \"54s\"\n      }\n    ]\n  }\n}\n","code":429,"status":"Too Many Requests"}}
    at throwErrorIfNotOK (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/genai/dist/node/index.mjs:14072:30)
    at process.processTicksAndRejections (node:internal/process/task_queues:95:5)
    at async file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/genai/dist/node/index.mjs:13863:13
    at async Models.generateContentStream (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/genai/dist/node/index.mjs:15204:24)
    at async LoggingContentGenerator.generateContentStream (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/core/loggingContentGenerator.js:57:22)
    at async retryWithBackoff (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/utils/retry.js:62:20)
    at async GeminiChat.makeApiCallAndProcessStream (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/core/geminiChat.js:266:32)
    at async file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/core/geminiChat.js:214:40
    at async Turn.run (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/core/turn.js:59:30)
    at async GeminiClient.sendMessageStream (file:///usr/local/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/core/client.js:405:26) {
  status: 429
}

cc @richjames0

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[RissyRan]

Now, this seems strange: https://screenshot.googleplex.com/6UjKsrTvNQPHMxV :(

[parambole]

Thanks for making these changes. Left a few comments.

[parambole]

LGTM. Thank you for adding this variant.

[RissyRan]

Some minor comments. Thanks!