[Model] DeepSeek-V4/V4.1 VL: stream language-model weights instead of sorting the checkpoint

[gitbisector]

Purpose

The DeepSeek-V4 VL wrapper and the NVIDIA DeepSeek-V4.1 VL wrapper sort the whole mapped weight stream so that the language_model. group reaches the child loader contiguously. The NVIDIA children finalize fused expert weights inside load_weights, so the group must not be interleaved. sorted() materializes the iterator, so every tensor the loader yields is alive before the first copy:

• Loaders that yield materialized tensors (--safetensors-load-strategy eager, or --load-format fastsafetensors, which yields a separate tensor per weight): nearly the whole checkpoint is held at once (DeepSeek-V4.1-Flash is 475 GiB). [Bugfix][ROCm] DeepSeek-V4/V4.1: Hybrid FP8/MXFP4 Quant Routing Guard, Fused Shared Experts, Engram Scale Translation, and VLM Generator Streaming #58782 reports host OOM from the same sort on the AMD V4.1 wrapper.
• The default mmap views: this costs no extra anonymous memory. But every shard stays mapped for the whole load and name order jumps between shards. In our runs that meant ~20–30 GiB more mapped page cache per rank, more storage reads and a slower load (below).

Contiguity needs only a non-interleaved group, not a total order. This PR adds stream_language_model_first, which yields language_model.* as it arrives and defers the rest (vision tower, aligner), cloned in case the loader reuses its buffers. Both wrappers use it.

Not a duplicate: #58782 applies the same idea to the AMD V4.1 wrapper as one part of a ROCm fix, and could reuse this helper. This PR covers the V4 wrapper and the NVIDIA V4.1 wrapper. #57733 keeps the in-load finalize, so the contiguity requirement stays.

Test Plan

pytest -v tests/models/test_deepseek_v4_vl_stream_weights.py

The test is parametrized over both wrappers and uses the real AutoWeightsLoader with a language model that finalizes in load_weights. It asserts that a weight is loaded before the stream is drained, that the child is called once, and that it finalizes after all of its weights. It is registered in the Basic Models (Other) CUDA step.

Test Result

Both cases pass on current main (GB10, torch 2.13.0+cu130). Without the change, both cases fail at the mid-stream assertion: the stream is drained before anything loads.

DeepSeek-V4.1-Flash at TP4 on 4x DGX Spark with --load-format auto, page cache dropped before each boot. These runs used a 0.28.1-based build with the same change minus the clone. Two nodes had read_ahead_kb=1024 and two had 128, the same in both arms, which is why reads vary by rank:

	boots	read per rank	load, slowest rank
sorted	1	179–357 GiB	248 s
streamed	2	158–265 GiB	202 s, 212 s

With --load-format fastsafetensors, #55985 and some fastsafetensors fixes that are not upstream yet, streaming is what lets this model load on a 121.7 GiB node at all: ~88 s, against 202–212 s for auto above.

This changes only the order in which weights reach the loader. The child's per-tensor loaders don't depend on order, and it still finalizes once, after all of its weights.

AI assistance (Claude Code) was used for the change, the tests and the measurements. I reviewed every line and ran the tests and boots.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gitbisector.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gitbisector.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork