Stream safetensors checkpoints into models

[gtong-nv]

Summary

• stream local or cached safetensors tensors directly into an already
  materialized model instead of first constructing a complete in-memory state
  dict
• use the safetensors file-backed loader for ordinary local safetensors loads
• validate checkpoint keys and tensor shapes before copying any weights, and
  periodically release consumed checkpoint pages from the Linux page cache
• add an opt-in stream_checkpoint setting to the WAN transformer config and
  enable it for every LingBot World v1 and v2 preset
• add CPU coverage for file-backed loading, direct-to-model streaming, and the
  LingBot configuration defaults

Motivation

Loading a large model previously required the destination model and the entire
checkpoint state dict to be resident at the same time. LingBot World uses a 14B
parameter transformer, so BF16 weights alone are roughly 28 GB before accounting
for the rest of the pipeline, allocator overhead, filesystem page cache, and
other runtime memory.

That transient duplication can exhaust system memory on machines with 64 GB of
host RAM. It is also a problem on systems with 128 GB of unified memory, such as
DGX Spark, because CPU checkpoint pages and GPU model allocations compete for
the same physical memory pool. The resulting memory pressure can cause an OOM,
heavy reclaim, or a load failure before inference starts even when the final
materialized model fits.

This change keeps checkpoint loading bounded by copying one tensor at a time
from the memory-mapped safetensors file into the destination model. It avoids a
second model-sized state dict and advises the OS to discard already consumed
checkpoint pages in 512 MiB intervals.

Conceptually, peak checkpoint-loading residency changes from approximately:

materialized model + complete checkpoint state dict + runtime overhead

to:

materialized model + current checkpoint tensor/mapped pages + runtime overhead

Behavior and compatibility

• Streaming is used when load_checkpoint(..., model=model) can resolve a
  local or cached safetensors file.
• LingBot World v1 and v2 explicitly request streaming through their shared
  base configuration.
• Existing WAN-based integrations retain the previous behavior unless they
  opt in with stream_checkpoint=True.
• Calls that request a state dict without supplying a model still return a
  state dict, preserving the public API.
• Non-safetensors checkpoints and checkpoint sources that cannot yet be
  resolved to a local file continue through the existing loader path.
• Streaming currently does not support state_dict_transform; the WAN loader
  raises a clear error if both options are requested.
• This affects model initialization only. It does not alter inference kernels
  or steady-state generation performance after loading completes.

Validation

• Added ci_cpu tests that reject the old whole-file byte-loading path for
  local safetensors.
• Added a direct-to-model test that verifies the complete state dict loader is
  not called and that the destination weights match the checkpoint.
• Added LingBot smoke coverage requiring all shipped v1 and v2 presets to
  enable streaming checkpoint loading.
• Ran the repository pre-commit lint checks.

Limitations and follow-ups

• A non-cached S3 safetensors checkpoint still falls back to the existing
  materializing path on its first load.
• A sharded safetensors checkpoint can stream from its merged local cache;
  producing that cache for the first time still uses the existing merge path.
• Supporting streaming key transforms would require a tensor-by-tensor remap
  interface rather than the current whole-state-dict callback.

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

[greptile-apps]

Greptile Summary

This PR adds file-backed safetensors loading and direct tensor-by-tensor streaming into materialized models, then enables the bounded-memory path for LingBot World presets.

• Validates checkpoint keys and shapes before copying weights.
• Adds page-cache eviction during large streamed loads.
• Adds a WAN configuration switch and LingBot defaults.
• Adds CPU coverage for file-backed and direct-to-model loading.

Confidence Score: 4/5

The dependency constraint must be aligned with the new safetensors API before merging, otherwise valid installations can fail during model initialization.

The streaming path calls safe_open with a keyword unavailable in safetensors releases still permitted by the project metadata, and shipped LingBot presets now reach that path.

Files Needing Attention: flashdreams/flashdreams/core/checkpoint/load.py and flashdreams/pyproject.toml

Important Files Changed

Filename	Overview
flashdreams/flashdreams/core/checkpoint/load.py	Adds file-backed and direct-to-model safetensors loading, but uses an API introduced after the package's declared minimum safetensors version.
flashdreams/flashdreams/recipes/wan/transformer/wan21.py	Adds an opt-in streaming configuration and preserves the existing transformed state-dict path when streaming is disabled.
integrations/lingbot/lingbot/config.py	Enables streaming for the LingBot base configuration and inherited presets, exposing the dependency-version incompatibility during cached checkpoint startup.
flashdreams/tests/test_checkpoint_loading.py	Covers file-backed and direct model loading but does not exercise the minimum supported safetensors version.
integrations/lingbot/tests/test_smoke.py	Verifies every shipped LingBot runner inherits the streaming checkpoint setting.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[WAN transformer initialization] --> B{stream_checkpoint enabled?}
  B -- No --> C[Materialize complete state dict]
  C --> D[Module.load_state_dict]
  B -- Yes --> E{Local or cached safetensors available?}
  E -- No --> C
  E -- Yes --> F[Open safetensors mmap]
  F --> G[Validate keys and shapes]
  G --> H[Copy one tensor at a time]
  H --> I[Periodically evict checkpoint pages]
  D --> J[Run post-load parameter updates]
  I --> J

Loading

_{Reviews (1): Last reviewed commit: "Stream safetensors checkpoints into mode..." | Re-trigger Greptile}

[greptile-apps]

Unsupported safetensors loader keyword

When an installation resolves safetensors 0.4 through 0.7 as permitted by the package metadata, this call passes an unsupported backend keyword, causing model initialization to raise TypeError for streamed LingBot and other direct-to-model checkpoint loads.

Suggested change

	with safe_open(path, framework="pt", device="cpu", backend="mmap") as source:
	with safe_open(path, framework="pt", device="cpu") as source:

Knowledge Base Used: Core Engine (flashdreams/flashdreams/core/)

[ArielG-NV]

/ok to test b0c0d2f

[ArielG-NV]

1 suggestion
1 question

[ArielG-NV]

If safetensors are available (can be automatically detected from a checkpoint file afaik), shouldn't we by default always load them?

[ArielG-NV]

We likely want this config option for all demos. Unsure if we should make this change now or if this change will be something we implement during a larger refactor.

[gtong-nv]

For smaller model or system with large RAM, this option is not helpful and will add some overhead when loading.
I'm testing the overhead to see if the difference matters or not. Will enable by default if not.