Cosmos3 super reasoner - #5074
Conversation
![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 adds support for NVIDIA Cosmos3-Nano and Cosmos3-Super Reasoner models, including their configurations, parameter mappings, and multimodal processing logic. It also introduces output truncation upon encountering an EOS token during decoding and switches safetensor loading to recursive search. The review feedback highlights a potential bug in handling list-valued eos_token_ids, a configuration mismatch regarding tied word embeddings for the Cosmos3-Nano model, and risks associated with recursively loading safetensors from subdirectories.
| eos_token_id = None | ||
| if hasattr(tokenizer_model, "tokenizer") and hasattr(tokenizer_model.tokenizer, "eos_token_id"): | ||
| eos_token_id = tokenizer_model.tokenizer.eos_token_id | ||
|
|
||
| max_logging.debug(f"using eos_token_id: {eos_token_id}") | ||
|
|
||
| if eos_token_id is not None and eos_token_id in results: | ||
| max_logging.info( | ||
| f"EOS token {eos_token_id} found at index {results.index(eos_token_id)};" | ||
| f" output sequence truncated to length {results.index(eos_token_id) + 1}" | ||
| ) | ||
| results = results[: results.index(eos_token_id) + 1] # Include the EOS token in the output |
There was a problem hiding this comment.
The current implementation assumes eos_token_id is always a single integer. However, in Hugging Face tokenizers (such as Gemma 4 or other models), eos_token_id can be a list or tuple of integers. If eos_token_id is a list, eos_token_id in results will evaluate to False (since results is a list of integers), and even if it matched, results.index(eos_token_id) would raise a ValueError because the list itself is not an element of results.
To prevent decoding truncation failures on models with multiple EOS tokens, normalize eos_token_id to a set of integers and find the first occurrence of any of those tokens.
| eos_token_id = None | |
| if hasattr(tokenizer_model, "tokenizer") and hasattr(tokenizer_model.tokenizer, "eos_token_id"): | |
| eos_token_id = tokenizer_model.tokenizer.eos_token_id | |
| max_logging.debug(f"using eos_token_id: {eos_token_id}") | |
| if eos_token_id is not None and eos_token_id in results: | |
| max_logging.info( | |
| f"EOS token {eos_token_id} found at index {results.index(eos_token_id)};" | |
| f" output sequence truncated to length {results.index(eos_token_id) + 1}" | |
| ) | |
| results = results[: results.index(eos_token_id) + 1] # Include the EOS token in the output | |
| eos_token_id = None | |
| if hasattr(tokenizer_model, "tokenizer") and hasattr(tokenizer_model.tokenizer, "eos_token_id"): | |
| eos_token_id = tokenizer_model.tokenizer.eos_token_id | |
| max_logging.debug(f"using eos_token_id: {eos_token_id}") | |
| if eos_token_id is not None: | |
| eos_tokens = {eos_token_id} if isinstance(eos_token_id, int) else set(eos_token_id) | |
| eos_indices = [results.index(t) for t in eos_tokens if t in results] | |
| if eos_indices: | |
| first_eos_idx = min(eos_indices) | |
| max_logging.info( | |
| f"EOS token found at index {first_eos_idx};" | |
| f" output sequence truncated to length {first_eos_idx + 1}" | |
| ) | |
| results = results[: first_eos_idx + 1] # Include the EOS token in the output |
| vocab_size: 151936 | ||
| normalization_layer_epsilon: 1.0e-6 | ||
| use_qk_norm: true | ||
| logits_via_embedding: false |
There was a problem hiding this comment.
There is a mismatch between the Hugging Face model configuration and the MaxText model configuration for the Cosmos3 Nano Reasoner. In cosmos3_nano_reasoner_dict (defined in hf_model_configs.py), "tie_word_embeddings" is set to True. However, in cosmos3-nano-reasoner.yml, logits_via_embedding is set to false.
To ensure consistency and avoid potential issues during checkpoint conversion or training, logits_via_embedding should be set to true to match the tied word embeddings configuration.
logits_via_embedding: true| ) | ||
| # load safetensors | ||
| ckpt_paths = sorted(pathlib.Path(local_path).glob("[!.]*.safetensors")) | ||
| ckpt_paths = sorted(pathlib.Path(local_path).rglob("[!.]*.safetensors")) |
There was a problem hiding this comment.
Changing .glob("[!.]*.safetensors") to .rglob("[!.]*.safetensors") recursively searches all subdirectories. If the local directory contains nested folders with other checkpoints (e.g., checkpoint-1000/model.safetensors alongside model.safetensors at the root), rglob will load all of them into hf_state_dict. This can lead to unexpected weight overwrites or out-of-memory errors.
Consider restricting the search to the root directory unless nested safetensors are explicitly expected, or filter out subdirectories that represent separate checkpoints.
c231d6c to
16453a7
Compare
1 participant
Description
[WIP] It contains messy manual rebase to the unsubmitted cosmos-nano PR. Will clean up once the base is merged.
This PR adds support for onboarding and running inference on the NVIDIA cosmos3-super.
Key updates
Reuse: It leverages the existing architectural blocks derived from Qwen3/Qwen3-VL, sharing preprocessing and attention mechanisms with the previously implemented Cosmos3-Nano variants while scaling up to a larger qwen3-vl backbone.
New components:
Fix the processor for cosmos models: the identifier should be the vision_block name rather than model name; thus the redundant explicit listing of cosmos3-xxx are removed.
Tests
Checkpoint conversion
Decoding (Robotic Arm Planning) (v6 lite)
Decoding output and analysis:
Input `<|im_start|>user <|vision_start|><|image_pad|><|vision_end|>You are a robotic arm planner. The task is to put the flower into the red bottle. Generate a detailed plan consisting of sequential subtasks to accomplish the task.<|im_end|> <|im_start|>assistant ` -> `Task Analysis:
The goal is to place the flower into the red bottle. To achieve this, we need to break down the task into a series of logical steps that the robotic arm can execute sequentially. Below is a detailed plan with explanations for each step.
Step 1: Identify and Locate the Flower
Step 2: Move the Robotic Arm to the Flower
Step 3: Grasp the Flower
Step 4: Lift the Flower
Step 5: Identify and Locate the Red Bottle
Step 6: Move the Robotic Arm to the Red Bottle
Step 7: Insert the Flower into the Red Bottle
Step 8: Release the Flower
Step 9: Verify the Task Completion
Final Answer:
The detailed plan to put the flower into the red bottle is as follows:
This sequence ensures that the task is completed accurately and efficiently.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.