Core 模型加载

安全 checkpoint 读取、state dict 发现、结构身份、惰性 DiskMap 与模型构造。

最后更新: 2026年7月22日

本页内容

一个安全的本地 state dict 示例如何选择加载层级 Checkpoint 信任边界完整参考

模型加载被分成多层，让调用方决定 Core 应当负责多少策略。load_torch_checkpoint 默认以 weights-only 安全模式读取一个 PyTorch 对象；load_state_dict 理解文件、目录、多路径、safetensors 和分片 index；DiskMap 按参数名惰性暴露 checkpoint tensor；load_model 还会构造 module、转换 key、分配权重、安装可选显存管理、移动 placement 并切换 eval 模式。

一个安全的本地 state dict 示例

from pathlib import Path
from tempfile import TemporaryDirectory

import torch

from worldfoundry.core import hash_state_dict_keys, load_torch_state_dict

with TemporaryDirectory() as directory:
    checkpoint = Path(directory) / "weights.pt"
    expected = {"linear.weight": torch.arange(8).reshape(2, 4)}
    torch.save(expected, checkpoint)

    loaded = load_torch_state_dict(checkpoint, map_location="cpu")
    assert torch.equal(loaded["linear.weight"], expected["linear.weight"])
    print(hash_state_dict_keys(loaded, with_shape=True))

hash_state_dict_keys 是参数名及可选形状的路由指纹，它有意忽略 tensor 值，因此架构相同的两个 checkpoint 可以得到相同 digest。需要确认文件内容身份时，应使用读取字节的 hash_model_file。

如何选择加载层级

已知文件就是 PyTorch checkpoint，并且需要保留原始外层结构时，用 load_torch_checkpoint。输入可能是目录、分片 index、safetensors 或多个路径，并且希望得到一个合并 mapping 时，用 load_state_dict。一次加载所有 tensor 会超过 host 内存时，用 DiskMap；safetensors 能提供最佳惰性行为，而二进制文件会走内存兼容 reader。

只有当共享构造路径符合模型时才使用 load_model。它会在 meta device 初始化上下文中构造，支持 state dict converter、DeepSpeed ZeRO-3 专属分配路径，并能按 module_map 安装 wrapper。参数物化方式特殊的模型应在自己的 runner 中控制这一层，并复用更低层的 Core 函数。

Checkpoint 信任边界

load_torch_checkpoint 默认使用 weights_only=True。可选的 allow_unsafe_pickle_fallback=True 可能执行 pickle payload，绝不能对不可信文件启用。Safetensors 没有这类 pickle 执行风险。远程 URI 会先经过 Core storage helper 本地化，再交给 reader。

完整参考

以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。

20 个公开符号

`assign_state_dict_strict`

func

def assign_state_dict_strict(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> Any

worldfoundry.core.assign_state_dict_strictfrom worldfoundry.core import assign_state_dict_strict

源码

简介

assign_state_dict_strict — Validate and assign tensors, including into a meta-device module. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：Any。

参数

moduleAny
state_dictMapping[str, Any]
labelstr: 默认值: 'checkpoint'

返回值: Any

`build_rename_dict`

func

def build_rename_dict(source_state_dict,target_state_dict,split_qkv = False)

worldfoundry.core.build_rename_dictfrom worldfoundry.core import build_rename_dict

源码

简介

build_rename_dict — Print parameter-key matches between two state dicts for conversion scripts. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

source_state_dict
target_state_dict
split_qkv: 默认值: False

`hash_model_file`

func

def hash_model_file(path, with_shape = True)

worldfoundry.core.hash_model_filefrom worldfoundry.core import hash_model_file

源码

简介

hash_model_file — Return an MD5 digest of checkpoint key names loaded from *path*. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

path
with_shape: 默认值: True

`hash_state_dict_keys`

func

def hash_state_dict_keys(state_dict, with_shape = True)

worldfoundry.core.hash_state_dict_keysfrom worldfoundry.core import hash_state_dict_keys

源码

简介

hash_state_dict_keys — Return a deterministic fingerprint of state-dict structure. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

state_dict: Possibly nested mapping whose tensor keys identify a model checkpoint layout.
with_shape: Include tensor dimensions as well as parameter names.默认值: True

说明

This is an identity hint, not a content or security hash: tensor values are not read. Use `hash_model_file` when file content integrity is required.

`load_model`

func

def load_model(model_class,path,config = None,torch_dtype = torch.bfloat16,device = 'cpu',state_dict_converter = None,use_disk_map = False,module_map = None,vram_config = None,vram_limit = None,state_dict = None)

worldfoundry.core.load_modelfrom worldfoundry.core import load_model

源码

简介

load_model — Construct a model, assign checkpoint weights, and finalize inference placement. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

model_class: PyTorch module class to instantiate.
path: Checkpoint path consumed by `load_state_dict or DiskMap`.
config: Keyword arguments passed to `model_class`.默认值: None
torch_dtype: Final model dtype.默认值: torch.bfloat16
device: Final model device when fine-grained VRAM management is absent.默认值: 'cpu'
state_dict_converter: Optional key/shape converter applied before assignment.默认值: None
use_disk_map: Read parameter tensors lazily instead of loading a full state dict.默认值: False
module_map: Source-class to VRAM-wrapper mapping. Enabling it delegates placement to `enable_vram_management`.默认值: None
vram_config: Offload/onload/preparing/computation placement dictionary.默认值: None
vram_limit: Optional used-memory limit in GiB for wrappers.默认值: None
state_dict: Already-loaded weights; takes precedence over `path`.默认值: None

说明

Construction uses meta-device initialization where possible. DeepSpeed ZeRO-3 receives its specialized state-dict assignment path.

`load_model_loader_registry`

func

def load_model_loader_registry(config_path: str | Path,model_classes: Mapping[str, type]) -> ModelLoaderRegistry

worldfoundry.core.load_model_loader_registryfrom worldfoundry.core import load_model_loader_registry

源码

简介

load_model_loader_registry — Load checkpoint-hash model routing rules from a package data YAML file. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：ModelLoaderRegistry。

参数

config_pathstr | Path
model_classesMapping[str, type]

返回值: ModelLoaderRegistry

`load_model_with_disk_offload`

func

def load_model_with_disk_offload(model_class,path,config = None,torch_dtype = torch.bfloat16,device = 'cpu',state_dict_converter = None,module_map = None)

worldfoundry.core.load_model_with_disk_offloadfrom worldfoundry.core import load_model_with_disk_offload

源码

简介

load_model_with_disk_offload — Construct a model whose inactive weights remain disk-backed. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

model_class: PyTorch module class to instantiate on the meta device.
path: Checkpoint path or paths indexed by `DiskMap`.
config: Keyword arguments passed to `model_class`.默认值: None
torch_dtype: Computation dtype.默认值: torch.bfloat16
device: Preparing and computation device.默认值: 'cpu'
state_dict_converter: Optional checkpoint-key converter.默认值: None
module_map: Required source-class to wrapper-class mapping.默认值: None

`load_safetensors_into_model_streaming`

func

def load_safetensors_into_model_streaming(model: Any,checkpoint_dir: str | os.PathLike[str],strict: bool = True,device: str | torch.device = 'cpu',dtype: torch.dtype | None = None,assign: bool | None = None) -> dict[str, int]

worldfoundry.core.load_safetensors_into_model_streamingfrom worldfoundry.core import load_safetensors_into_model_streaming

源码

简介

load_safetensors_into_model_streaming — Load checkpoint shards one at a time and validate the complete key set. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：dict[str, int]。

源码 docstring

Load checkpoint shards one at a time and validate the complete key set.

Unlike a merged state-dict loader, peak host memory is bounded by the largest shard. Meta-initialized models are materialized tensor-by-tensor directly on `device; floating checkpoint tensors may also be converted to dtype` while streaming. Strict validation is performed after every shard has been applied, so the result is equivalent to one strict full-checkpoint load.

参数

modelAny
checkpoint_dirstr | os.PathLike[str]
strictbool: 默认值: True
devicestr | torch.device: 默认值: 'cpu'
dtypetorch.dtype | None: 默认值: None
assignbool | None: 默认值: None

返回值: dict[str, int]

`load_sharded_safetensors_parallel_with_progress`

func

def load_sharded_safetensors_parallel_with_progress(checkpoint_dir: str)

worldfoundry.core.load_sharded_safetensors_parallel_with_progressfrom worldfoundry.core import load_sharded_safetensors_parallel_with_progress

源码

简介

load_sharded_safetensors_parallel_with_progress — Load a safetensors checkpoint, reading independent shards concurrently. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

checkpoint_dirstr: Directory containing either `model.safetensors or a model.safetensors.index.json plus its shards. A sibling .zst file is decompressed through the system zstd` command.

异常

FileNotFoundError: Neither the index nor fallback model file exists.
RuntimeError: External zstd decompression fails.

`load_state_dict`

func

def load_state_dict(file_path,torch_dtype = None,device = 'cpu',pin_memory = False,verbose = 0)

worldfoundry.core.load_state_dictfrom worldfoundry.core import load_state_dict

源码

简介

按 Core 的放置与 key 处理策略，把权重载入模块。

参数

file_path: Local/remote checkpoint URI, directory, safetensors index, or list of any of those. Later files overwrite duplicate keys.
torch_dtype: Optional dtype conversion applied to tensor values.默认值: None
device: Device used while deserializing weights; CPU is the safe default.默认值: 'cpu'
pin_memory: Pin CPU tensors after loading to accelerate a later GPU copy.默认值: False
verbose: Print start/finish messages when at least `1`.默认值: 0

说明

File type is selected from the path: directories are scanned, `*.safetensors.index.json follows shards, *.safetensors` uses safetensors, and other suffixes use the safe PyTorch checkpoint loader.

`load_state_dict_from_safetensors_index`

func

def load_state_dict_from_safetensors_index(file_path,torch_dtype = None,device = 'cpu')

worldfoundry.core.load_state_dict_from_safetensors_indexfrom worldfoundry.core import load_state_dict_from_safetensors_index

源码

简介

load_state_dict_from_safetensors_index — Load every unique shard referenced by a safetensors index. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

file_path: Local or remote `*.safetensors.index.json` URI.
torch_dtype: Optional dtype conversion for each tensor.默认值: None
device: Device used by the safetensors reader.默认值: 'cpu'

异常

ValueError: The index has no mapping-valued `weight_map`.

`load_torch_checkpoint`

func

def load_torch_checkpoint(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu',weights_only: bool | None = True,allow_unsafe_pickle_fallback: bool = False,**kwargs: Any) -> Any

worldfoundry.core.load_torch_checkpointfrom worldfoundry.core import load_torch_checkpoint

源码

简介

load_torch_checkpoint — Load a PyTorch checkpoint with weights-only deserialization by default. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：Any。

参数

checkpoint_pathstr | os.PathLike[str]: Local or remote checkpoint URI.
map_locationAny: Destination understood by `torch.load`.默认值: 'cpu'
weights_onlybool | None: Safe deserialization mode. `None` omits the argument for compatibility with older PyTorch releases.默认值: True
allow_unsafe_pickle_fallbackbool: Retry with unrestricted pickle only after a weights-only unpickling failure. Never enable for untrusted files.默认值: False
kwargsAny: Additional `torch.load` options.

警告

Setting `allow_unsafe_pickle_fallback=True` can execute code embedded in a malicious checkpoint.

返回值: Any — Object produced by `torch.load`.

`load_torch_state_dict`

func

def load_torch_state_dict(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu') -> Any

worldfoundry.core.load_torch_state_dictfrom worldfoundry.core import load_torch_state_dict

源码

简介

load_torch_state_dict — Load a PyTorch state-dict-shaped checkpoint in weights-only mode. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：Any。

参数

checkpoint_pathstr | os.PathLike[str]: Local or remote checkpoint URI.
map_locationAny: Destination understood by `torch.load`.默认值: 'cpu'

返回值: Any — Deserialized checkpoint object. The function does not unwrap outer `state_dict/module keys; use load_state_dict` for that policy.

`ModelConfig`

cls

class ModelConfig(path: Union[str, list[str]] = None,model_id: str = None,origin_file_pattern: Union[str, list[str]] = None,download_source: str = None,local_model_path: str = None,skip_download: bool = None,offload_device: Optional[Union[str, torch.device]] = None,offload_dtype: Optional[torch.dtype] = None,onload_device: Optional[Union[str, torch.device]] = None,onload_dtype: Optional[torch.dtype] = None,preparing_device: Optional[Union[str, torch.device]] = None,preparing_dtype: Optional[torch.dtype] = None,computation_device: Optional[Union[str, torch.device]] = None,computation_dtype: Optional[torch.dtype] = None,clear_parameters: bool = False,state_dict: Dict[str, torch.Tensor] = None)

worldfoundry.core.ModelConfigfrom worldfoundry.core import ModelConfig

源码

简介

ModelConfig — Resolved model source, download policy, and VRAM placement settings. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

属性

pathUnion[str, list[str]]: Existing checkpoint path or paths. When set, no model-hub lookup is needed.默认值: None
model_idstr: Hugging Face or ModelScope repository identifier.默认值: None
origin_file_patternUnion[str, list[str]]: Optional allow-pattern within the model repository.默认值: None
download_sourcestr: `"huggingface" or "modelscope"; defaults from WORLDFOUNDRY_DOWNLOAD_SOURCE` and then to Hugging Face.默认值: None
local_model_pathstr: Local repository cache root. Defaults through `WORLDFOUNDRY_MODEL_DIR`.默认值: None
skip_downloadbool: Reuse local files without contacting the model hub.默认值: None
offload_deviceOptional[Union[str, torch.device]]: Device used while weights are inactive.默认值: None
offload_dtypeOptional[torch.dtype]: Dtype used while weights are inactive.默认值: None
onload_deviceOptional[Union[str, torch.device]]: First-stage prefetch device.默认值: None
onload_dtypeOptional[torch.dtype]: First-stage prefetch dtype.默认值: None
preparing_deviceOptional[Union[str, torch.device]]: Second-stage prefetch device.默认值: None
preparing_dtypeOptional[torch.dtype]: Second-stage prefetch dtype.默认值: None
computation_deviceOptional[Union[str, torch.device]]: Device used for forward execution.默认值: None
computation_dtypeOptional[torch.dtype]: Dtype used for forward execution.默认值: None
clear_parametersbool: Compatibility flag for loaders that release source tensors after assignment.默认值: False
state_dictDict[str, torch.Tensor]: Optional already-loaded weights, bypassing file loading.默认值: None

方法

methcheck_input()源码

简介

check_input — Require either a concrete path or a model-hub identifier.

methparse_original_file_pattern()源码

简介

parse_original_file_pattern — Normalize the repository allow-pattern to a glob-compatible value.

methparse_download_source()源码

简介

parse_download_source — Resolve the configured model hub, including the environment override.

methparse_skip_download()源码

简介

parse_skip_download — Resolve offline behavior from the field or environment.

methdownload()源码

简介

download — Download missing repository files into `local_model_path`.

methrequire_downloading(use_usp: bool = False)源码

简介

require_downloading — Return whether this rank should contact the configured model hub.

参数

use_uspbool: 默认值: False

methreset_local_model_path()源码

简介

reset_local_model_path — Apply the canonical WorldFoundry model directory when needed.

methdownload_if_necessary(use_usp: bool = False)源码

简介

download_if_necessary — Materialize the configured source and replace `path` with local files.

源码 docstring

Materialize the configured source and replace `path` with local files.

In USP execution only rank zero downloads; all ranks synchronize before the final local path is resolved.

参数

use_uspbool: 默认值: False

methvram_config()源码

简介

vram_config — Return placement fields in the mapping expected by VRAM wrappers.

`ModelLoaderRegistry`

cls

class ModelLoaderRegistry(model_loader_configs: list[tuple[Any, str, list[str], list[type], str]],huggingface_model_loader_configs: list[tuple[str, str, str, Any]],patch_model_loader_configs: list[tuple[str, list[str], list[type], dict[str, Any]]],preset_models_on_huggingface: dict[str, Any],preset_models_on_modelscope: dict[str, Any],preset_model_ids: tuple[str, ...],preset_model_websites: tuple[str, ...])

worldfoundry.core.ModelLoaderRegistryfrom worldfoundry.core import ModelLoaderRegistry

源码

简介

ModelLoaderRegistry — Validated routing data used to choose a model loader from checkpoint identity. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

属性

model_loader_configslist[tuple[Any, str, list[str], list[type], str]]: Single-file hash rules and their resolved model classes.
huggingface_model_loader_configslist[tuple[str, str, str, Any]]: Architecture-to-library routing rules.
patch_model_loader_configslist[tuple[str, list[str], list[type], dict[str, Any]]]: Patch/adaptor hash rules and extra kwargs.
preset_models_on_huggingfacedict[str, Any]: Named Hugging Face model presets.
preset_models_on_modelscopedict[str, Any]: Named ModelScope model presets.
preset_model_idstuple[str, ...]: Stable preset identifiers.
preset_model_websitestuple[str, ...]: Human-facing model source labels.

`search_for_embeddings`

func

def search_for_embeddings(state_dict)

worldfoundry.core.search_for_embeddingsfrom worldfoundry.core import search_for_embeddings

源码

简介

search_for_embeddings — Return all tensor leaves from a nested state dict. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

state_dict

`search_for_files`

func

def search_for_files(folder, extensions)

worldfoundry.core.search_for_filesfrom worldfoundry.core import search_for_files

源码

简介

search_for_files — Recursively find files matching any suffix in `extensions`. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

folder
extensions

`search_parameter`

func

def search_parameter(param,state_dict,atol = 0.001)

worldfoundry.core.search_parameterfrom worldfoundry.core import search_parameter

源码

简介

search_parameter — Find the first state-dict key whose tensor numerically matches `param`. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

param
state_dict
atol: 默认值: 0.001

`split_state_dict_with_prefix`

func

def split_state_dict_with_prefix(state_dict)

worldfoundry.core.split_state_dict_with_prefixfrom worldfoundry.core import split_state_dict_with_prefix

源码

简介

split_state_dict_with_prefix — Split a state dict by the first dotted parameter-key segment. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。

参数

state_dict: Mapping with string parameter names.

`validate_state_dict_compatibility`

func

def validate_state_dict_compatibility(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> None

worldfoundry.core.validate_state_dict_compatibilityfrom worldfoundry.core import validate_state_dict_compatibility

源码

简介

validate_state_dict_compatibility — Reject missing, unexpected, or shape-incompatible checkpoint tensors. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。标注返回类型：None。

源码 docstring

Reject missing, unexpected, or shape-incompatible checkpoint tensors.

This validation is intentionally independent of tensor dtype: released FP32 weights may be assigned to a meta-device model and converted to the selected inference dtype during final placement.

参数

moduleAny
state_dictMapping[str, Any]
labelstr: 默认值: 'checkpoint'

返回值: None