[HuiyingLi]

What does this PR do ?

Add native AutoModel support for deepseek-ai/DeepSeek-V4-Flash-Vision-Exp, including the released vision tower/aligner, dynamic-resolution processor, image-token bridge, visual attention visibility, vision-aware MoE routing, checkpoint conversion, FSDP/PP integration, and MedPix validation recipes.

Changelog

• Implement the DeepSeek V4 vision transformer with 2D RoPE, fp32 RMSNorm, full patch attention, SwiGLU blocks, and the spatial downsampling aligner.
• Add the dynamic-resolution image processor and released N-layout pseudo-token expansion.
• Merge encoded image features and sentinel embeddings into the DSV4 language stream on the first pipeline stage.
• Match released image-span visibility for dense and TileLang sparse attention.
• Route visual tokens with bias_vl while preserving text hash routing in the early DSV4 layers.
• Map vision, aligner, sentinel, and visual-routing checkpoint tensors through the state-dict adapter.
• Add FSDP/PP layer discovery and fp32 norm-island handling for the vision tower.
• Add the 16-node MedPix 4K validation recipe and focused unit/parity tools.

Validation

Vision-only released-implementation parity

This parity compares only the 32-layer vision encoder, aligner, visual N-layout, and permutation using the released checkpoint implementation and the same real checkpoint weights.

Slurm job 17526208: COMPLETED, exit 0:0, 1 GPU, 41 seconds.

Source image (W x H)	ViT grid (H x W)	Vision max abs	Aligner max abs	N-layout	Permutation
384 x 384	28 x 28	0	0	exact	exact
511 x 377	27 x 37	0	0	exact	exact
377 x 511	37 x 27	0	0	exact	exact
1000 x 60	10 x 78	0	0	exact	exact
60 x 1000	112 x 7	0	0	exact	exact

Focused tests

• python -m pytest -q tests/unit_tests/models/deepseek_v4/test_dsv4_vision.py
• Result: 60 passed on one H100.
• Covers dynamic image sizes, N-layout/permutation, ragged grids and backward, image-span visibility, sparse-window indices and dense-vs-TileLang forward/backward parity, visual MoE routing, prompt processing, checkpoint mapping, PP media staging, layer discovery, and FSDP vision norm policy.

Distributed training validation

• Single-node gate, Slurm 17527387: COMPLETED 0:0; 8 GPUs; full 32-layer vision tower; 4 language layers; 4K sequence; EP4/PP2; HybridEP + TileLang + torch-mm experts; activation checkpointing; two backward/optimizer steps.
• Scaled EP32 smoke, Slurm 17529267: COMPLETED 0:0; 8 nodes / 64 GPUs; 4K sequence; EP32; two backward/optimizer steps.
• Final MedPix run, Slurm 17529854: COMPLETED 0:0; 16 nodes / 128 GPUs; full model; 4K sequence; EP32/PP4; HybridEP + TileLang + torch-mm experts; activation checkpointing; 100 steps.

Final 100-step run: Weights & Biases

• Dataset: mmoukouba/MedPix-VQA, train split
• Loss: 2.2812 at step 0 to 1.4531 at step 99
• Mean loss: first 10 steps 2.12265; last 10 steps 1.42031
• All losses and gradient norms finite
• Steady-state time: 11.19 s/step
• Step 99 throughput: 3273.26 tokens/s (102.29 tokens/s per EP/data-parallel GPU)
• Step 99 peak memory: 63.89 GiB

Existing DSV4 CP/Tulu-3 regression

The existing text-only DeepSeek-V4 Flash training path was validated on the exact PR head with the NeMo-CI nightly image. No DSV4 CP/Tulu-3 regression was observed.

• Root pipeline 65561789: success
• Training job 420080099: success; Slurm 5956453, COMPLETED, exit 0:0
• Source: PR head 9cbdddc1dbbeeaeb6a7a1e15ca50d36cc1a9cab7, runtime-mounted over nvcr.io/nvidian/nemo-automodel:nightly; no container rebuild
• Topology: 16 nodes / 128 H100; FSDP2; TP1/CP8/PP4/EP32; HybridEP dispatcher; TileLang attention; torch-mm experts; activation checkpointing
• Data path: allenai/tulu-3-sft-mixture, train split; 16K sequence; max_length padding; unpacked; global batch 256 / local batch 4
• Completed 10/10 forward, backward, and optimizer steps
• Losses: 1.1989, 1.1164, 0.9533, 0.7001, 0.7561, 0.8824, 0.9414, 0.6268, 0.6302, 0.7628
• All losses and gradient norms were finite; steady-state peak memory was 35.49 GiB/GPU

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?

Additional Information

• The visual path intentionally uses unpacked CP1 execution. Packed THD/context-parallel visual spans are not added by this PR.
• Visual MTP token shifting is not defined by the released path, so the provided visual fine-tuning recipes set num_nextn_predict_layers: 0.

[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.

[HuiyingLi]

/claude review

[claude]

Reviewed the DeepSeek-V4-Flash-Vision-Exp onboarding (16 files; uv.lock excluded). This is a large but carefully structured and well-tested change. I confirmed the aspects most likely to break:

• Registry & capabilities: DeepseekV4ForCausalLM is already in MODEL_ARCH_MAPPING and its ModelCapabilities already declares supports_pp/ep/cp=True, so the new pp4×ep32 recipe does not introduce capability staleness.
• Checkpoint-free init: the vision gate's router weight is initialized via the base Gate branch in MoE.init_weights, and init_dsv4_weights() initializes bias_vl/tid2eid; sentinels and vision/aligner params are initialized in initialize_weights. Routing indices from tid2eid are in-range and topk-distinct.
• Image-visibility mask under PP: both vision_token_types and attention_mask are batch-major kwargs that ride schedule.step and are re-derived/rebuilt per stage, so the bidirectional span mask is applied on every PP stage (media pixels are correctly kept off the schedule via stage_vlm_media_for_pp).
• fp32 self-casting FSDP policy and the bidirectional state-dict adapter renames (vision./aligner./sentinels/bias_vl) round-trip and are covered by tests.
• Sparse-window kernel rewrite matches the released reference oracle and does not introduce a trainable-parameter nondeterministic-accumulation hazard.

One minor tensor-contract nit left inline (missing layout docstring for the new vision_token_types argument on build_dsv4_sparse_topk_indices). Nothing blocking.

[claude]

This signature was materially changed to accept a new tensor input vision_token_types, but the one-line docstring doesn't document its layout. The body validates vision_token_types.shape == (batch_size, seq_len), so the contract should be stated. Suggest expanding to a Google-style Args section, e.g.:

vision_token_types: Optional pseudo-token types of shape [batch, sequence]
    (`-1` at text positions, `IMAGE_START`/`IMAGE`/`IMAGE_END` etc. inside
    image spans). When provided with `max_image_tokens > 0`, image queries
    additionally see the full image span; ignored otherwise.
max_image_tokens: Maximum tokens per image span; `0` disables visual windowing.

[HuiyingLi]

/claude review

[claude]

LGTM

Reviewed all 16 changed files (uv.lock excluded). This adds the DeepSeek-V4-Flash-Vision tower, aligner, visual MoE gate, image-visibility masking, sparse visual attention windowing, processor, PP media staging, state-dict mapping, and parity tooling. Verified against the trusted base:

• Registry/capabilities already cover DSV4 (supports_pp/ep/cp=True); the new PP4×EP32 YAML introduces no capability staleness.
• DeepseekV4VisionGate subclasses Gate and overrides only _route_scores, preserving the shared aux-loss/forward path and fp32 selection-bias buffers.
• Checkpoint-free init overwrites every sentinel torch.empty parameter and the tid2eid/bias_vl buffers; the deterministic hash map yields in-range, per-token-distinct expert IDs for the DeepEP top-k contract.
• fp32 RoPE/RMSNorm islands are preserved: initialize_weights casts via cast_model_to_dtype, and the new preserve_activation_dtype FSDP policy keeps bf16 activations (tested).
• Image-visibility mask and visual sparse-window indices are transcriptions of the released reference, gated to unpacked CP=1, and covered by oracle + forward/backward parity tests.
• PP media reuses the established _vlm_*_chunks side channel; tensor-contract docstrings document layouts; state-dict round-trip is tested; tool scripts use weights_only=True.

[HuiyingLi]

/ok to test fd0b017