Skip to content

Add Qwen3 1.7B configs#3103

Open
rxu183 wants to merge 5 commits intoAI-Hypercomputer:mainfrom
rxu183:richard/qwen1_7b
Open

Add Qwen3 1.7B configs#3103
rxu183 wants to merge 5 commits intoAI-Hypercomputer:mainfrom
rxu183:richard/qwen1_7b

Conversation

@rxu183
Copy link

@rxu183 rxu183 commented Feb 6, 2026 •
edited
Loading

Description

Adds Qwen3 1.7B model configs. It's extremely similar to other Qwen model configs, with slight changes to the base_emb_dim and base_mlp_dim relative to the 0.6B. I've made a slight change to the documentation listing 1.7B as a supported model.

Tests

I tested this on a Google Colab v6e1 instance, via the provided Qwen SFT demo notebook. The training cell with TFLOPs and loss succeeded, so I'm pretty confident the architecture mapping and parameter conversion was done properly. To reproduce, all I changed was the notebook to use 1.7B instead of the 0.6B model. However, I wasn't able to run the cells regarding vLLM, due to a ipython issue.

vLLM Error 

---------------------------------------------------------------------------
UnsupportedOperation                      Traceback (most recent call last)
[/tmp/ipython-input-2349923095.py](https://localhost:8080/#) in <cell line: 0>()
      4 
      5 tunix_model = TunixMaxTextAdapter(trainer.model)
----> 6 vllm_rollout = VllmRollout(
      7     model=tunix_model,
      8     tokenizer=tokenizer,

17 frames
[/usr/local/lib/python3.12/dist-packages/ipykernel/iostream.py](https://localhost:8080/#) in fileno(self)
    309             return self._original_stdstream_copy
    310         else:
--> 311             raise io.UnsupportedOperation("fileno")
    312 
    313     def _watch_pipe_fd(self):

UnsupportedOperation: fileno
Logs 

INFO:absl:Lazy loading DISABLED. Loading full HuggingFace model: Qwen/Qwen3-1.7B...
config.json: 100% 726/726 [00:00<00:00, 4.57MB/s]
model.safetensors.index.json: 25.6kB [00:00, 118MB/s]
Fetching 2 files: 0% 0/2 [00:00<?, ?it/s]
model-00002-of-00002.safetensors: 0% 0.00/622M [00:00<?, ?B/s]

...

Starting SFT Training...
Training: 100%
 500/500 [01:22<00:00, 12.23step/s, _train_loss=0.764, _train_perplexity=2.15, _train_steps_per_sec=12.1]
Per train step:
 Total TFLOPs: 10.93 
 split as 96.70% learnable weight flops and 3.30% attention flops
SFT Training Complete!

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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@shralex
Copy link
Collaborator

shralex commented Feb 10, 2026

Thank you for the contribution. Can you please rebase ?

@RissyRan
Copy link
Collaborator

RissyRan commented Feb 10, 2026

Thanks for bringup new models! We usually run this script to test correctness against HF version. Please feel free to let us know if you meet any issues.

cc @parambole who is working on Qwen3 bringup and could help PR review.

hengtaoguo
hengtaoguo reviewed Feb 12, 2026
View reviewed changes
Copy link
Collaborator

@hengtaoguo hengtaoguo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally LGTM! Have you run forward_pass_logit_check for this 1.7B model?

https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/checkpointing_solutions/convert_checkpoint.md#verifying-conversion-correctness

@rxu183
Copy link
Author

rxu183 commented Feb 12, 2026

No, not yet unfortunately - I'm out of town and haven't had much of a chance to look at it quite yet, but will after I get back. And of course, happy to contribute!

@rxu183
Copy link
Author

rxu183 commented Feb 16, 2026 •
edited
Loading

Hi, thanks for advising me about the script, and here are the logs that I got when I ran it. I think they're similar enough, but will let y'all be the judge of that.

Logs
Details 
INFO:absl:
--- Prompt: I love to ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1349       | read                 | 21.1250    |
| 4296       | cook                 | 20.7500    |
| 1486       | play                 | 20.5000    |
| 3270       | write                | 20.3750    |
| 3960       | learn                | 20.3750    |
| 5821       | travel               | 20.3750    |
| 8180       | eat                  | 20.2500    |
| 1490       | see                  | 20.1250    |
| 1281       | make                 | 20.0000    |
| 990        | use                  | 19.8750    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1349       | read                 | 21.1250    |
| 4296       | cook                 | 20.7500    |
| 1486       | play                 | 20.5000    |
| 3960       | learn                | 20.3750    |
| 3270       | write                | 20.3750    |
| 5821       | travel               | 20.3750    |
| 8180       | eat                  | 20.2500    |
| 1490       | see                  | 20.1250    |
| 1281       | make                 | 20.0000    |
| 387        | be                   | 19.8750    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 9/10                 |
| jaccard_similarity             | 0.8181818181818182   |
| rank_agreement_percentage      | 70.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.002645
INFO:absl:
Max KL divergence for a single token in the set: 0.003262
INFO:absl:
--- Prompt: Today is a ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1899       | day                  | 21.1250    |
| 3281       | special              | 20.6250    |
| 7014       | Monday               | 20.5000    |
| 7728       | Saturday             | 19.7500    |
| 7589       | Tuesday              | 19.6250    |
| 6233       | beautiful            | 19.6250    |
| 7920       | Wednesday            | 19.5000    |
| 7270       | Sunday               | 19.3750    |
| 1602       | very                 | 19.3750    |
| 7794       | Thursday             | 19.2500    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1899       | day                  | 21.2500    |
| 3281       | special              | 20.6250    |
| 7014       | Monday               | 20.5000    |
| 6233       | beautiful            | 19.7500    |
| 7728       | Saturday             | 19.7500    |
| 7920       | Wednesday            | 19.5000    |
| 7589       | Tuesday              | 19.5000    |
| 1602       | very                 | 19.3750    |
| 7270       | Sunday               | 19.3750    |
| 7794       | Thursday             | 19.2500    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 40.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.001330
INFO:absl:
Max KL divergence for a single token in the set: 0.003074
INFO:absl:
--- Prompt: What is the ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 7428       | purpose              | 22.3750    |
| 1887       | main                 | 22.3750    |
| 6672       | difference           | 22.3750    |
| 25361      | significance         | 21.8750    |
| 5025       | relationship         | 21.8750    |
| 7271       | definition           | 21.8750    |
| 14806      | formula              | 21.6250    |
| 897        | value                | 21.5000    |
| 3491       | problem              | 21.3750    |
| 3476       | role                 | 21.2500    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1887       | main                 | 22.2500    |
| 7428       | purpose              | 22.2500    |
| 6672       | difference           | 22.2500    |
| 5025       | relationship         | 21.8750    |
| 7271       | definition           | 21.7500    |
| 25361      | significance         | 21.7500    |
| 14806      | formula              | 21.5000    |
| 897        | value                | 21.3750    |
| 3476       | role                 | 21.2500    |
| 3491       | problem              | 21.2500    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 30.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.002301
INFO:absl:
Max KL divergence for a single token in the set: 0.008017

@hengtaoguo
Copy link
Collaborator

hengtaoguo commented Feb 18, 2026

Hi, thanks for advising me about the script, and here are the logs that I got when I ran it. I think they're similar enough, but will let y'all be the judge of that.

Logs
Details

INFO:absl:
--- Prompt: I love to ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1349       | read                 | 21.1250    |
| 4296       | cook                 | 20.7500    |
| 1486       | play                 | 20.5000    |
| 3270       | write                | 20.3750    |
| 3960       | learn                | 20.3750    |
| 5821       | travel               | 20.3750    |
| 8180       | eat                  | 20.2500    |
| 1490       | see                  | 20.1250    |
| 1281       | make                 | 20.0000    |
| 990        | use                  | 19.8750    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1349       | read                 | 21.1250    |
| 4296       | cook                 | 20.7500    |
| 1486       | play                 | 20.5000    |
| 3960       | learn                | 20.3750    |
| 3270       | write                | 20.3750    |
| 5821       | travel               | 20.3750    |
| 8180       | eat                  | 20.2500    |
| 1490       | see                  | 20.1250    |
| 1281       | make                 | 20.0000    |
| 387        | be                   | 19.8750    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 9/10                 |
| jaccard_similarity             | 0.8181818181818182   |
| rank_agreement_percentage      | 70.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.002645
INFO:absl:
Max KL divergence for a single token in the set: 0.003262
INFO:absl:
--- Prompt: Today is a ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1899       | day                  | 21.1250    |
| 3281       | special              | 20.6250    |
| 7014       | Monday               | 20.5000    |
| 7728       | Saturday             | 19.7500    |
| 7589       | Tuesday              | 19.6250    |
| 6233       | beautiful            | 19.6250    |
| 7920       | Wednesday            | 19.5000    |
| 7270       | Sunday               | 19.3750    |
| 1602       | very                 | 19.3750    |
| 7794       | Thursday             | 19.2500    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1899       | day                  | 21.2500    |
| 3281       | special              | 20.6250    |
| 7014       | Monday               | 20.5000    |
| 6233       | beautiful            | 19.7500    |
| 7728       | Saturday             | 19.7500    |
| 7920       | Wednesday            | 19.5000    |
| 7589       | Tuesday              | 19.5000    |
| 1602       | very                 | 19.3750    |
| 7270       | Sunday               | 19.3750    |
| 7794       | Thursday             | 19.2500    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 40.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.001330
INFO:absl:
Max KL divergence for a single token in the set: 0.003074
INFO:absl:
--- Prompt: What is the ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 7428       | purpose              | 22.3750    |
| 1887       | main                 | 22.3750    |
| 6672       | difference           | 22.3750    |
| 25361      | significance         | 21.8750    |
| 5025       | relationship         | 21.8750    |
| 7271       | definition           | 21.8750    |
| 14806      | formula              | 21.6250    |
| 897        | value                | 21.5000    |
| 3491       | problem              | 21.3750    |
| 3476       | role                 | 21.2500    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1887       | main                 | 22.2500    |
| 7428       | purpose              | 22.2500    |
| 6672       | difference           | 22.2500    |
| 5025       | relationship         | 21.8750    |
| 7271       | definition           | 21.7500    |
| 25361      | significance         | 21.7500    |
| 14806      | formula              | 21.5000    |
| 897        | value                | 21.3750    |
| 3476       | role                 | 21.2500    |
| 3491       | problem              | 21.2500    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 30.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 0.002301
INFO:absl:
Max KL divergence for a single token in the set: 0.008017

This looks great, Max KL divergence is below the threshold 0.015 which verifies the correctness.

@codecov
Copy link

codecov bot commented Feb 18, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@hengtaoguo hengtaoguo hengtaoguo approved these changes

@jacoguzo jacoguzo Awaiting requested review from jacoguzo jacoguzo is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@rxu183 @shralex @RissyRan @hengtaoguo

Comments