Skip to content

Parameterize bos_id in inference_microbenchmark.py#3786

Merged
copybara-service[bot] merged 1 commit intomainfrom
fix-bos-token
May 1, 2026
Merged

Parameterize bos_id in inference_microbenchmark.py#3786
copybara-service[bot] merged 1 commit intomainfrom
fix-bos-token

Conversation

@A9isha
Copy link
Copy Markdown
Collaborator

@A9isha A9isha commented Apr 30, 2026

Description

Problem:
inference_microbenchmark.py crashes when running models without a BOS token (like Qwen3) because it hardcodes is_bos=True during tokenization. This causes JetStream's HuggingFaceTokenizer to prepend None to the token list, resulting in a JAX TypeError.

Solution:
Dynamically set is_bos = tokenizer_model.bos_id is not None to avoid prepending None for models that don't have a BOS token.

FIXES: b/508334388

Tests

Ran the following command successfully on a v5p-8

python3 -m maxtext.inference.inference_microbenchmark   model_name=qwen3-4b   tokenizer_path=Qwen/Qwen3-4B   per_device_batch_size=1   max_prefill_predict_length=1024   max_target_length=2048   ici_autoregressive_parallelism=4   scan_layers=false   weight_dtype=bfloat16   attention=dot_product

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.

Copy link
Copy Markdown
Collaborator

@parambole parambole left a comment

Choose a reason for hiding this comment

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

LGTM 🔥

@copybara-service copybara-service Bot merged commit bb2e6d5 into main May 1, 2026
46 of 48 checks passed
@copybara-service copybara-service Bot deleted the fix-bos-token branch May 1, 2026 17:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants