feat: Add gemma4 drafter model support

[athitten]

What does this PR do ?

This PR adds joint fine-tuning support for Gemma 4 base (also called target model) and drafter/assistant models that enable multi-token prediction (MTP). The drafter is co-trained with the Gemma 4 base end-to-end via a composite model (Gemma4WithDrafter) that wires up shared K/V states, sqrt(H_b)-scaled embeddings, and a K-step recurrent forward matching the Gemma 4 drafter tech report. The PR provides two reference configs for joint fine-tuning of gemma-4-E4B-it and gemma-4-E4B-it-assistant, one with MedPix VQA dataset and the other with a text-only Tulu-3 + Magicoder mix. Also provides an inference benchmark script that validates speculative-decode throughput on the saved checkpoint.

The feature has been verified only against the Gemma 4 4B (E4B) base + drafter pair. The composite is architecturally model-agnostic within the Gemma 4 family, but the example YAMLs, the parity tests, and all training/inference verification in this PR target the 4B pair.

Implementation

• Gemma4WithDrafter composite (nemo_automodel/components/models/gemma4_drafter/)

  • Wraps a Gemma4ForConditionalGeneration base + Gemma4AssistantForCausalLM drafter as a single nn.Module for FSDP2 training.
  • K-step recurrent forward: round 0 feeds cat(base.embed(input_ids), base.h_final); round k≥1 feeds cat(base.embed(input_ids_shifted_by_k), prev_drafter.last_hidden_state). shared_kv_states is captured once from a single base forward and reused at every round (per the Gemma 4 drafter tech report).
  • Saves to <ckpt>/base/ and <ckpt>/drafter/ HF-loadable subdirs for vLLM / HF inference handoff.
  • Conditionally unfreezes post_projection only when drafter_num_steps > 1
  • masked_embedding.centroids always frozen (the torch.topk inside it blocks gradient flow back to the centroids).

• Recipe (recipes/vlm/finetune.py, recipes/base_recipe.py)

  • New _maybe_add_drafter_loss(out, base_loss, labels, …) sums per-step CE over the composite's drafter_logits list with _shift_labels_left(labels, k) per round.
  • base_recipe.load_checkpoint dispatches to model.load_pretrained when the composite exposes it, so the base/ + drafter/ subdir layout reloads correctly on resume.

• Text-only dataset adapter (components/datasets/vlm/datasets.py)

  • make_tulu3_magicoder_text_mix_dataset interleaves allenai/tulu-3-sft-mixture (80 %) and ise-uiuc/Magicoder-OSS-Instruct-75K (20 %) into {"conversation": [...]} dicts with no image field, so default_collate_fn emits batches without pixel_values. The composite + base accept text-only inputs unchanged.

• Example configs and benchmark (examples/vlm_finetune/gemma4_joint_drafter/)

  • gemma4_4b_joint_drafter_medpix.yaml — Joint fine-tuning recipe of google/gemma-4-E4B-it + google/gemma-4-E4B-it-assistant with MedPix VQA
  • gemma4_4b_joint_drafter_tulu_magicoder_mix.yaml — Joint fine-tuning on tulu and magiccoder mix, same base + drafter pair
  • benchmark_mtp_inference.py — measures speculative-decode acceptance + throughput end-to-end against the trained 4B pair

Testing and Verification

All numbers below are for joint fine-tuning of gemma-4-E4B-it with /gemma-4-E4B-it-assistant

1. Fine-tuning the joint model on MedPix-VQA: no NaNs, no loss spikes, no grad-norm spikes. Using the recipe examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_medpix.yaml
   Loss curve:

2. Large-scale fine-tuning on a Tulu-3 (80 %) + Magicoder (20 %) mix for 500 steps. Stable fine-tuning run observed with the recipe examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_tulu_magicoder_mix.yaml
   Loss curve:

3. Inference run on the tulu + magicoder fine-tuned checkpoint after 500 steps. Ran benchmark_mtp_inference.py against the saved base/ + drafter/ pair. Results below:

4. Verified save + standalone HF reload works correctly. Checkpoints are stored in separate <ckpt>/base/ and <ckpt>/drafter/ subdirs as a prerequisite for generation. After a 5-step training run, each sub-checkpoint loads via plain HF from_pretrained without any NeMo-specific code:

   • Gemma4ForConditionalGeneration.from_pretrained("<ckpt>/base/model/consolidated") → 7.94 B params ✓
   • Gemma4AssistantForCausalLM.from_pretrained("<ckpt>/drafter/model/consolidated") → 78.5 M params ✓
   • State-dict key sets match the upstream released references (google/gemma-4-E4B-it and google/gemma-4-E4B-it-assistant) exactly (0 missing, 0 extra).
   • masked_embedding.token_ordering (int64[262144]) buffer survives the DCP → safetensors path.

5. Verified resume-from-checkpoint loss parity. Run A: 10 steps fresh with ckpt_every_steps=5. Run B: same config but --checkpoint.restore_from <RunA>/epoch_0_step_4. Per-step loss is bit-identical between A and B at every overlapping step (5–9): 3.1257, 3.2992, 3.4045, 2.5375, 3.2824. Confirms model weights + optimizer state + LR scheduler state + dataloader state + RNG all restore correctly.

6. Unit and functional tests.

• tests/unit_tests/models/gemma4_drafter/test_composite.py — composite pre-projection input layout, shared-KV plumbing, K-step recurrence, save/load round-trip.
• tests/unit_tests/models/gemma4_drafter/test_composite_fsdp2.py — FSDP2 wrap + expert/grad sync.
• tests/unit_tests/models/gemma4_drafter/test_drafter_wrapper.py — drafter sub-module integration.
• tests/unit_tests/recipes/test_vlm_drafter_helpers.py — _shift_labels_left + _maybe_add_drafter_loss.
• tests/functional_tests/L2_HF_Transformer_VLM_Gemma4_Joint_Drafter.sh — L2 single-node smoke.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[athitten]

/ok to test 753ee70