Add multi-directory checkpoint loader for multimodal alignment poc

[subawocit]

Description

PR draft

Usage

Define hf_token and base_output_directory in generate_ckpt.py, then:

python3 -m maxtext.experimental.omni_poc.generate_ckpt

The checkpoints from stitched model will then be saved in your output directory.

Tests

Unit Test Results

================================================================================
TEST 1: MERGING AND ASSEMBLY CORRECTNESS (MERGING VISION AND LANGUAGE MODEL CHECKPOINTS)
================================================================================
  - Verified LLM subtrees restore from LLM checkpoint: Passed
  - Verified Vision subtrees restore from Vision checkpoint: Passed
  - Verified LLM layers count & MLP intermediate dimension consistency: Passed (orig=36, stitched=36) for LLM 'qwen3-4b'
  - Verified Vision layers count consistency: Passed (orig=27, stitched=27) for Vision 'gemma3-4b'
  - Overall Status: PASSED
================================================================================

================================================================================
TEST 2: OUTPUT LOGITS COMPARISON (ORIGINAL VS. STITCHED CHECKPOINTS)
================================================================================
  Raw Image Input       : shape=[2, 896, 896, 3]       dtype=float32
  Text Prompt Tokens    : shape=[2, 128]               dtype=int32
  Vision Tower Output   : Max Diff = 0.000000e+00
  Projector Output      : Max Diff = 11290753.0000
  Token Embedder Output : Max Diff = 0.000000e+00
================================================================================

================================================================================
TEST 3: MULTIMODAL DATA FLOW & MIDDLE SHAPES
================================================================================
  Stage 1: Raw Image Input    : shape=[2, 896, 896, 3]       dtype=float32
  Stage 2: Vision Tower Output: shape=[2, 1, 256, 1152]      dtype=float32 (1 image, 256 soft tokens, dim 1152)
  Stage 3: Stitched Projector Output: shape=[2, 256, 2560]         dtype=float32 (Projected & squeezed to LLM dim 2560)
  Stage 4: Text Prompt Tokens : shape=[2, 128]               dtype=int32 (128 text token IDs)
  Stage 5: Text Token Embeddings: shape=[2, 128, 2560]         dtype=float32 (Text embeddings, dim 2560)
================================================================================

  - Overall Status: PASSED

Step-by-Step Objectives

• ☑️ Multi-Directory Checkpoint Restoration
• Dynamic MLP Connector & LLM Freezing
• Special Tokenizer & Placeholder Masking
• COCO-Narratives Data Pipeline Processing
• SFT Execution & Evaluation

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!