Cosmos3 super reasoner

[lydhr]

Description

This PR adds support for onboarding and running inference on the NVIDIA cosmos3-super, reasoner tower.

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:

  • Dimensions: $d_{model}=5120$, $L_{heads}=64$, $L_{kv_heads}=8$.
  • Capacity: 64 decoder layers.

• 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

MODEL_NAME=cosmos3-super-reasoner

PYTHONPATH=src/ python3 -m maxtext.checkpoint_conversion.to_maxtext \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    base_output_directory=<Your GS bucket or local>\
    scan_layers=false \
    use_multimodal=true \
    hardware=cpu \
    skip_jax_distributed_system=True \
    checkpoint_storage_use_ocdbt=False \
    checkpoint_storage_use_zarr3=False \
    --eager_load_method=safetensors \
    --lazy_load_tensors=False

Decoding (Robotic Arm Planning) (v6 lite)

MODEL_NAME=cosmos3-super-reasoner
export checkpoint_path=gs://ly-maxtext/output/cosmos3-super-reasoner/09031507/0/items

PYTHONPATH=src/ python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    tokenizer_path=nvidia/Cosmos3-Super \
    tokenizer_type=huggingface  \
    load_parameters_path=${checkpoint_path} \
    per_device_batch_size=1 \
    run_name=${MODEL_NAME?}_decode \
    scan_layers=false \
    use_multimodal=true \
    prompt='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.' \
    image_path='tests/assets/test_image_reasoning.jpg' \
    max_prefill_predict_length=1024 \
    max_target_length=2248 \. # increased the length to see the full reasoning result.
    ici_tensor_parallelism=4 \
    override_model_config=true \
    attention='dot_product'

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

• Objective: Determine the position of the flower on the table.
• Explanation: The flower is currently lying flat on the table. The robotic arm must first locate it using its sensors or vision system. This step ensures that the arm knows where to move to pick up the flower.

Step 2: Move the Robotic Arm to the Flower

• Objective: Position the robotic arm's gripper over the flower.
• Explanation: Once the flower's location is identified, the robotic arm needs to move its gripper to the exact coordinates of the flower. This involves precise movement along the X, Y, and Z axes to align the gripper directly above the flower.

Step 3: Grasp the Flower

• Objective: Securely pick up the flower using the gripper.
• Explanation: The gripper must close around the flower's stem to ensure a firm grip. It is important to avoid damaging the flower during this process. The gripper should be designed to handle delicate objects like flowers.

Step 4: Lift the Flower

• Objective: Raise the flower off the table.
• Explanation: After grasping the flower, the robotic arm lifts it vertically to clear the table surface. This step ensures that the flower is fully detached from the table and ready to be moved toward the red bottle.

Step 5: Identify and Locate the Red Bottle

• Objective: Determine the position of the red bottle on the table.
• Explanation: The red bottle is one of the three bottles on the table. The robotic arm must use its sensors or vision system to identify which bottle is red and note its coordinates.

Step 6: Move the Robotic Arm to the Red Bottle

• Objective: Position the robotic arm's gripper over the red bottle.
• Explanation: With the flower securely held, the robotic arm moves horizontally and vertically to bring the flower above the opening of the red bottle. Precision is crucial here to ensure the flower is aligned correctly for insertion.

Step 7: Insert the Flower into the Red Bottle

• Objective: Place the flower inside the red bottle.
• Explanation: The robotic arm carefully lowers the flower into the red bottle. The gripper should release the flower once it is fully inside the bottle. This step requires gentle control to avoid knocking over the bottle or damaging the flower.

Step 8: Release the Flower

• Objective: Ensure the flower is securely placed inside the red bottle.
• Explanation: After inserting the flower, the gripper opens to release it. The robotic arm then retracts slightly to confirm that the flower is stable within the bottle.

Step 9: Verify the Task Completion

• Objective: Confirm that the flower is successfully placed in the red bottle.
• Explanation: The robotic arm may use its sensors or vision system to verify that the flower is indeed inside the red bottle. If any issues are detected (e.g., the flower is not properly placed), the arm can repeat the necessary steps.

---

Final Answer:

The detailed plan to put the flower into the red bottle is as follows:

1. Identify and Locate the Flower
2. Move the Robotic Arm to the Flower
3. Grasp the Flower
4. Lift the Flower
5. Identify and Locate the Red Bottle
6. Move the Robotic Arm to the Red Bottle
7. Insert the Flower into the Red Bottle
8. Release the Flower
9. Verify the Task Completion

By following these steps sequentially, the robotic arm will successfully complete the task of placing the flower into the red bottle.

1. Identify and Locate the Flower
2. Move the Robotic Arm to the Flower
3. Grasp the Flower
4. Lift the Flower
5. Identify and Locate the Red Bottle
6. Move the Robotic Arm to the Red Bottle
7. Insert the Flower into the Red Bottle
8. Release the Flower
9. Verify the Task Completion`

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.

[gemini-code-assist]

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!