Core 模型加载
安全 checkpoint 读取、state dict 发现、结构身份、惰性 DiskMap 与模型构造。
模型加载被分成多层，让调用方决定 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 个公开符号
def assign_state_dict_strict(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> Anyassign_state_dict_strict — Validate and assign tensors, including into a meta-device module. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。 标注返回类型：Any。
参数
moduleAnystate_dictMapping[str, Any]labelstr默认值 'checkpoint'返回值
def build_rename_dict(source_state_dict,target_state_dict,split_qkv = False)build_rename_dict — Print parameter-key matches between two state dicts for conversion scripts. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
source_state_dicttarget_state_dictsplit_qkv默认值 Falsedef hash_model_file(path,with_shape = True)hash_model_file — Return an MD5 digest of checkpoint key names loaded from *path*. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
pathwith_shape默认值 Truedef hash_state_dict_keys(state_dict,with_shape = True)hash_state_dict_keys — Return a deterministic fingerprint of state-dict structure. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
state_dictwith_shape默认值 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.
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,strict = True,post_load_hook = None)load_model — Construct a model, assign checkpoint weights, and finalize inference placement. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
model_classpathconfig默认值 Nonemodel_class`.torch_dtype默认值 torch.bfloat16device默认值 'cpu'state_dict_converter默认值 Noneuse_disk_map默认值 Falsemodule_map默认值 Noneenable_vram_management`.vram_config默认值 Nonevram_limit默认值 Nonestate_dict默认值 Nonepath`.strict默认值 Trueinitialize(missing_keys)` when the model provides it.post_load_hook默认值 None返回值
说明
Construction uses meta-device initialization where possible. DeepSpeed ZeRO-3 receives its specialized state-dict assignment path.
def load_model_loader_registry(config_path: str | Path,model_classes: Mapping[str, type]) -> ModelLoaderRegistryload_model_loader_registry — Load checkpoint-hash model routing rules from a package data YAML file. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。 标注返回类型：ModelLoaderRegistry。
参数
config_pathstr | Pathmodel_classesMapping[str, type]def load_model_with_disk_offload(model_class,path,config = None,torch_dtype = torch.bfloat16,device = 'cpu',state_dict_converter = None,module_map = None)load_model_with_disk_offload — Construct a model whose inactive weights remain disk-backed. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
model_classpathDiskMap`.config默认值 Nonemodel_class`.torch_dtype默认值 torch.bfloat16device默认值 'cpu'state_dict_converter默认值 Nonemodule_map默认值 None返回值
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]简介
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.
参数
modelAnycheckpoint_dirstr | os.PathLike[str]strictbool默认值 Truedevicestr | torch.device默认值 'cpu'dtypetorch.dtype | None默认值 Noneassignbool | None默认值 None返回值
def load_sharded_safetensors_parallel_with_progress(checkpoint_dir: str)load_sharded_safetensors_parallel_with_progress — Load a safetensors checkpoint, reading independent shards concurrently. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
checkpoint_dirstrmodel.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.
def load_state_dict(file_path,torch_dtype = None,device = 'cpu',pin_memory = False,verbose = 0)按 Core 的放置与 key 处理策略，把权重载入模块。
参数
file_pathtorch_dtype默认值 Nonedevice默认值 'cpu'pin_memory默认值 Falseverbose默认值 01`.返回值
说明
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.
def load_state_dict_from_safetensors_index(file_path,torch_dtype = None,device = 'cpu')load_state_dict_from_safetensors_index — Load every unique shard referenced by a safetensors index. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
file_path*.safetensors.index.json` URI.torch_dtype默认值 Nonedevice默认值 'cpu'返回值
异常
ValueError- The index has no mapping-valued `
weight_map`.
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) -> Anyload_torch_checkpoint — Load a PyTorch checkpoint with weights-only deserialization by default. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。 标注返回类型：Any。
参数
checkpoint_pathstr | os.PathLike[str]map_locationAny默认值 'cpu'torch.load`.weights_onlybool | None默认值 TrueNone` omits the argument for compatibility with older PyTorch releases.allow_unsafe_pickle_fallbackbool默认值 FalsekwargsAnytorch.load` options.返回值
torch.load`.警告
Setting `allow_unsafe_pickle_fallback=True` can execute code embedded in a malicious checkpoint.
def load_torch_state_dict(checkpoint_path: str | os.PathLike[str],map_location: Any = 'cpu') -> Anyload_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]map_locationAny默认值 'cpu'torch.load`.返回值
state_dict/module keys; use load_state_dict` for that policy.class ModelConfig(path: Optional[Union[str, list[str]]] = None,model_id: Optional[str] = None,origin_file_pattern: Optional[Union[str, list[str]]] = None,download_source: Optional[str] = None,local_model_path: Optional[str] = None,skip_download: Optional[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: Optional[Dict[str, torch.Tensor]] = None)ModelConfig — Resolved model source, download policy, and VRAM placement settings. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
属性
pathOptional[Union[str, list[str]]]默认值 Nonemodel_idOptional[str]默认值 Noneorigin_file_patternOptional[Union[str, list[str]]]默认值 Nonedownload_sourceOptional[str]默认值 None"huggingface" or "modelscope"; defaults from WORLDFOUNDRY_DOWNLOAD_SOURCE` and then to Hugging Face.local_model_pathOptional[str]默认值 NoneWORLDFOUNDRY_MODEL_DIR`.skip_downloadOptional[bool]默认值 Noneoffload_deviceOptional[Union[str, torch.device]]默认值 Noneoffload_dtypeOptional[torch.dtype]默认值 Noneonload_deviceOptional[Union[str, torch.device]]默认值 Noneonload_dtypeOptional[torch.dtype]默认值 Nonepreparing_deviceOptional[Union[str, torch.device]]默认值 Nonepreparing_dtypeOptional[torch.dtype]默认值 Nonecomputation_deviceOptional[Union[str, torch.device]]默认值 Nonecomputation_dtypeOptional[torch.dtype]默认值 Noneclear_parametersbool默认值 Falsestate_dictOptional[Dict[str, torch.Tensor]]默认值 None方法
check_input — Require either a concrete path or a model-hub identifier.
parse_original_file_pattern — Normalize the repository allow-pattern to a glob-compatible value.
parse_download_source — Resolve the configured model hub, including the environment override.
简介
parse_skip_download — Resolve offline behavior from the field or environment.
源码 docstring
Resolve offline behavior from the field or environment.
Accepts the common boolean spellings; an unparseable value raises instead of silently falling through to "download anyway", which would hang offline clusters whose users set e.g. `=1` expecting a skip.
返回值
download — Download missing repository files into `local_model_path`.
require_downloading — Return whether this rank should contact the configured model hub.
参数
use_uspbool默认值 False简介
reset_local_model_path — Apply the canonical WorldFoundry model directory when none is set.
源码 docstring
Apply the canonical WorldFoundry model directory when none is set.
An explicitly configured `local_model_path always wins; WORLDFOUNDRY_MODEL_DIR is consumed by local_model_root_path()` and only fills the default (explicit argument > environment > default).
简介
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默认值 Falsevram_config — Return placement fields in the mapping expected by VRAM wrappers.
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, ...])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]]huggingface_model_loader_configslist[tuple[str, str, str, Any]]patch_model_loader_configslist[tuple[str, list[str], list[type], dict[str, Any]]]preset_models_on_huggingfacedict[str, Any]preset_models_on_modelscopedict[str, Any]preset_model_idstuple[str, ...]preset_model_websitestuple[str, ...]def search_for_embeddings(state_dict)search_for_embeddings — Return all tensor leaves from a nested state dict. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
state_dictdef search_for_files(folder,extensions)search_for_files — Recursively find files matching any suffix in `extensions`. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
folderextensionsdef search_parameter(param,state_dict,atol = 0.001)search_parameter — Find the first state-dict key whose tensor numerically matches `param`. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
paramstate_dictatol默认值 0.001def split_state_dict_with_prefix(state_dict)split_state_dict_with_prefix — Split a state dict by the first dotted parameter-key segment. 属于 Core 模型加载（checkpoint、state dict、DiskMap、构造）。
参数
state_dict返回值
def validate_state_dict_compatibility(module: Any,state_dict: Mapping[str, Any],label: str = 'checkpoint') -> None简介
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.
参数
moduleAnystate_dictMapping[str, Any]labelstr默认值 'checkpoint'