Stream safetensors checkpoints into models

[gtong-nv]

Summary

• Stream LingBot’s sharded safetensors checkpoint directly into the materialized model.
• Skip creation of the duplicate merged safetensors cache for streamed loads.
• Retain the state-dict loading path for non-safetensors checkpoints and models that require whole-state-dict transforms.
• Enable streaming by default for LingBot presets.

Resource savings

LingBot’s 16 transformer shards total approximately 69 GiB. Previously, non-streaming loads also created a merged safetensors cache, requiring an additional ~70 GiB of disk space.

Streaming avoids this duplicate cache and avoids retaining a complete checkpoint state dict during loading. This reduces peak host-memory pressure and prevents out-of-memory failures on RAM-constrained systems.

Loading-time delta

Mode	Ready time (2 runs)	Median
Non-streaming	20.708 s, 20.306 s	20.507 s
Streaming	23.576 s, 23.694 s	23.635 s

Benchmarked on an NVIDIA GB300 with cached LingBot checkpoints, streaming adds 3.13 s (15.3%) to startup time. This ~3-second cost is an acceptable tradeoff for the 14B model’s reduced disk usage and host-memory pressure.

stream_checkpoint is enabled by default for LingBot and remains disabled by default for other models until their checkpoint layouts and transform requirements are validated.

[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

The PR adds file-backed and direct-to-model safetensors loading to reduce peak checkpoint memory, then enables it for LingBot presets.

• Streams single-file and sharded safetensors tensors into materialized models.
• Validates checkpoint keys and tensor shapes before copying weights.
• Adds configurable WAN initialization devices and LingBot streaming defaults.
• Updates LingBot WebRTC argument aliases, runtime device propagation, and loading tests.

Confidence Score: 4/5

The PR is not yet safe to merge because supported safetensors installations can still fail during streamed model initialization.

The streaming paths pass backend="mmap" to safe_open, while the package metadata continues to permit safetensors versions predating that keyword, so model initialization can raise TypeError before loading any weights.

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

Important Files Changed

Filename	Overview
flashdreams/flashdreams/core/checkpoint/load.py	Adds file-backed loading and single-file/sharded direct-to-model streaming, but the previously reported safetensors version-contract mismatch remains.
flashdreams/flashdreams/recipes/wan/transformer/wan21.py	Adds opt-in streaming checkpoint loading and configurable initial parameter allocation on the runtime device.
integrations/lingbot/lingbot/config.py	Enables bounded streaming checkpoint loading for the shared LingBot configuration and its derived presets.
integrations/lingbot/lingbot/webrtc/session.py	Propagates the selected runtime device into WAN transformer initialization.
integrations/lingbot/lingbot/webrtc/server.py	Adds kebab-case aliases for several CLI options and derives the default preset from the runtime configuration.
flashdreams/tests/test_checkpoint_loading.py	Adds CPU tests for file-backed local loading and direct streaming of single-file and sharded checkpoints.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[LingBot or WAN transformer config] --> B[Materialize model on init_device]
  B --> C{stream_checkpoint enabled?}
  C -- No --> D[Load complete state dict]
  C -- Yes --> E{Checkpoint format}
  E -- Single safetensors --> F[Resolve local or cached file]
  E -- Sharded index --> G[Resolve index and shards]
  F --> H[Validate keys and shapes]
  G --> H
  H --> I[Copy one tensor at a time]
  I --> J[Initialized model]

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into dev/gtong/strea..." | Re-trigger Greptile}

[ArielG-NV]

/ok to test b0c0d2f

[ArielG-NV]

1 suggestion
1 question

[gtong-nv]

Mode	Ready time (2 runs)	Median
Non-streaming	20.708 s, 20.306 s	20.507 s
Streaming	23.576 s, 23.694 s	23.635 s

Benchmarked on GB300, lingbot checkpoints streaming is 3.13 s / 15.3% slower on cached startup
I think we should still keep streaming enabled for LingBot. It avoids the non-streaming path’s additional 70 GB merged safetensors cache and host-residency risk; the ~3 s startup cost is a reasonable tradeoff for this 14B model.

[gtong-nv]

/ok to test 2aaeb74

[ArielG-NV]

Benchmarked on GB300, lingbot checkpoints streaming is 3.13 s / 15.3% slower on cached startup
I think we should still keep streaming enabled for LingBot. It avoids the non-streaming path’s additional 70 GB merged safetensors cache and host-residency risk; the ~3 s startup cost is a reasonable tradeoff for this 14B model.

Is this because we are issuing 1 big copy (_copy_checkpoint_tensor), because tensor_names is very large (code below), or something-else (a collection of costs)?

                for name in tensor_names:
                    tensor = source.get_tensor(name)
                    try:
                        tensor_bytes = _copy_checkpoint_tensor(
                            model_state[name], tensor
                        )
                        shard_copied_bytes += tensor_bytes
                        total_copied_bytes += tensor_bytes
                    finally:
                        del tensor

Is this something we can potentially improve upon in your opinion?