cosmos3_nano model bring up (patch PR#4475)

[lydhr]

Description

This PR adds support for the nvidia/Cosmos3-Nano model in MaxText.

Key Changes

• Parameter Mapping: Implemented specialized COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING.
  • Instead of maintaining duplicate specialized mappings and hook functions for Cosmos3 in PR#4475, this PR reuses QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING and QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN.
  • A lightweight translation helper is appended to handle Cosmos3-specific naming conventions (like self_attn.to_q), improving maintainability.
• Vision Integration: Configured to use the qwen3_vl vision encoder block.

Reference

• Patched and adapted from an early workstream #4475.

---

Tests

Checkpoint Conversion

MODEL_NAME=cosmos3-nano-reasoner
python3 -m maxtext.checkpoint_conversion.to_maxtext \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    base_output_directory=${MODEL_BUCKET} \
    scan_layers=false \
    use_multimodal=true \
    hf_access_token=${HF_TOKEN?} \
    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 / Inference (Robotic Arm Planning)

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=${MODEL_NAME?} \
    tokenizer_path=nvidia/Cosmos3-Nano \
    tokenizer_type=huggingface  \
    load_parameters_path=${checkpoint_path} \
    per_device_batch_size=1 \
    run_name=sft_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=1124 \
    ici_tensor_parallelism=4 \
    override_model_config=true \
    attention='dot_product' \
    hf_access_token=${HF_TOKEN}

Sample Output:

->Move the arm to the flower
Move the arm to the red bottle
Place the flower in the red bottle
Move the arm away from the red bottle
Move the arm to the center of the table
Move the arm to the left side of the table
Move the arm to the right side of`

---

Checklist

• 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 introduces support for the NVIDIA Cosmos3-Nano Reasoner model, adding its configuration, parameter mapping, and integration across multimodal processing, rotary embeddings, and decoders. However, several critical issues must be addressed before merging. Specifically, there are multiple unresolved git merge conflict markers left in param_mapping.py and types.py. Additionally, the newly added unit tests in param_mapping_test.py reference a non-existent function (COSMOS3_TEXT_MAXTEXT_TO_HF_PARAM_MAPPING) and need to explicitly set maxtext_config.use_multimodal = False to avoid raising a KeyError during execution.

[codecov]

Codecov Report

❌ Patch coverage is 61.19403% with 52 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/utils/param_mapping.py	64.22%	40 Missing and 4 partials ⚠️
src/maxtext/multimodal/processor.py	14.28%	4 Missing and 2 partials ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	1 Missing ⚠️
src/maxtext/layers/attentions.py	0.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!